NewType for Callables in Python

Question

I would like to define two types Shape and Square. Here Square is a subclass of Shape.

Shape = NewType('Shape', tuple[int, int])
Square = NewType('Square', Shape)

def square_area(sh: Square) -> int:
    a, b = sh
    return a * b

def shape_area(sh: Shape) -> int:
    a, b = sh
    return a * b

I would define an additional type SquareCalculator as follows:

SquareCalculator = NewType('Calculator', Callable[[Square], int])

def get_calc() -> SquareCalculator:
    return square_area

The above code leads to the following errors in mypy:

newtyp.py:10: error: String argument 1 "Calculator" to NewType(...) does not match variable name "SquareCalculator"
newtyp.py:13: error: Incompatible return value type (got "Callable[[Square], int]", expected "SquareCalculator")
Found 2 errors in 1 file (checked 1 source file)

Is there a way around this? Note that I would like to maintain the difference in Shape and Square to ensure type-checking catches errors when they are inter-changed.

Answer 1

SquareCalculator should be a type alias, not a newtype, otherwise you need to use typing.cast.

Current syntax:

from typing import TypeAlias
from collections.abc import Callable

SquareCalculator: TypeAlias = Callable[[Square], int]

Python 3.12 syntax:

from collections.abc import Callable

type SquareCalculator = Callable[[Square], int]