Defining a type where the runtime type is selected at runtime

Typing
1

I am trying to figure out the best way to implement type annotations for types that are selected dynamically at runtime.

I have a type called MPZ that represents an implementation of integers and is defined at runtime by something like:

try:
    from flint import fmpz
    MPZ = fmpz
except ImportError:
    MPZ = int

I then have various functions that either work with int or with MPZ e.g.:

def f(x: int) -> int:
    x_mpz = MPZ(x)
    result_mpz = g(x_mpz)
    return int(result_mpz)

def g(x: MPZ) -> MPZ:
    return x + MPZ(1)

It is possible that at runtime MPZ and int are the same type but it is also possible that they are not. The int and fmpz types are sufficiently compatible that they can do what I need them to do but it is also important to distinguish between int and MPZ in the codebase because accepting/leaking an MPZ (possibly an fmpz) at runtime can cause problems elsewhere.

I would like for a typechecker to distinguish int and MPZ as incompatible types e.g.:

def f() -> MPZ:
    return 0 # error

def g() -> int:
    return MPZ(0) # error

Currently the best I have come up with for this is something like:

from __future__ import annotations
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    class MPZ:
        def __new__(cls, x: int) -> MPZ: ...
        def __int__(self) -> int: ...
        def __add__(self, other: MPZ) -> MPZ: ...
else:
    try:
        from flint import fmpz
        MPZ = fmpz
    except ImportError:
        MPZ = int

def f(x: int) -> int:
    return int(g(MPZ(x)))

def g(x: MPZ) -> MPZ:
    return x + MPZ(1)

print(f(1))

def h() -> int:
    return MPZ(1) # error

def i() -> MPZ:
    return 1 # error

This works as intended in all ways but one: it has a downside that much of the interesting checking is bypassed in the else clause of the if TYPE_CHECKING block. There are no stubs files for python-flint but for MPZ = int a type checker could verify the assignment against MPZ as e.g. a protocol but here all checking is disabled.

I have tried using protocols but then it means that I can’t use the MPZ type for annotations because it is a runtime variable:

from __future__ import annotations
from typing import Protocol, Self

class MPZ_proto(Protocol):
    def __new__(cls, x: int, /) -> Self: ...
    def __int__(self) -> int: ...
    def __add__(self, other: Self, /) -> Self: ...

MPZ: type[MPZ_proto]

try:
    from flint import fmpz # type: ignore
    MPZ = fmpz
except ImportError:
    MPZ = int

# Variable not allowed in type expression
def g(x: MPZ) -> MPZ: # error
    return x + MPZ(1)

Hopefully I am just missing something obvious about how to do this…

2

mypy and pyright both seem fine with this. It does mean explicitly having the difference between what is sufficiently “MPZ-like” separate from the actual implementation, but this allows the proper distinction to be expressible in places where it matters.

from __future__ import annotations
from typing import Protocol, Self

class MPZLike(Protocol):
    def __new__(cls, x: int, /) -> Self: ...
    def __int__(self) -> int: ...
    def __add__(self, other: Self, /) -> Self: ...


try:
    from flint import fmpz # type: ignore
    MPZ = fmpz
except ImportError:
    MPZ = int


def g(x: MPZLike) -> MPZLike:
    return x + MPZ(1)
1 Like
3

Thanks. This works for ensuring that a function with -> int only returns an int but not for ensuring that I have MPZ when I should rather than int. An example is:

def g(x: MPZLike) -> MPZLike:
    return x + 1

g(1) # checks fine

Now the function g is operating with int when I wanted it to use MPZ. I can also end up having a mixed list of MPZ and int:

a: list[MPZLike] = [1, 2, MPZ(3)]

For examples of why I want to care about this in both directions see e.g. broken int printing and perf problems with int or gmpy2.mpz broken with log. The situation here is that I can have a protocol to regulate my own usage but I still always have a strong preference about which type it is regardless of whether they seem to be compatible with a protocol: if I say it is MPZ then it should be MPZ and not int and vice versa.

4

Would bounded Generics work for this?

def g[MyInt: (int, fmpz)](x: MyInt) -> MyInt:
    return x + MPZ(1)

typing.TypeVar declarations can define bounds and constraints with a colon (: ) followed by an expression. A single expression after the colon indicates a bound (e.g. T: int ).

5

I don’t think that works or maybe I have misunderstood. The problem is that the type checker does not know the relationship between MyInt and MPZ:

from __future__ import annotations
from typing import Protocol, Self

class MPZLike(Protocol):
    def __new__(cls, x: int, /) -> Self: ...
    def __int__(self) -> int: ...
    def __add__(self, other: Self, /) -> Self: ...

MPZ: type[MPZLike]

try:
    from flint import fmpz # type: ignore
    MPZ = fmpz
except ImportError:
    MPZ = int

def g[MyInt: (int, fmpz)](x: MyInt) -> MyInt:
    # Operator + not supported for MyInt and MPZLike
    return x + MPZ(1) # error

g(1)

I need MPZ to be a runtime object that constructs instances but I also want it to be a type somehow that a type checker can understand. As far as I can tell the only way for a type checker to understand that MPZ is both a type and also a runtime callable object that constructs instances is if you write class MPZ but then it isn’t possible to have the necessary assignment MPZ = int without disabling the type checker somehow.

6

Oh I see. Thanks.

I can’t get it to typecheck, but it’s not possible to separate out all the times an instance of MPZ is created directly, and replace it with a call to a simple factory, that can be given one of two implementations in the same try block?

7

You should be using if TYPE_CHECKING in some form for sure - type checkers are not currently able to deal variables of any kind in annotations and a few of the implementers are resistant to adding this.

I think you can create a Protocol and a “concrete” class, the later being used in annotations, and the former being both the base class for the concrete class and used in an assert_type(int, MPZProtocol) to check that it is compatible.

8

Yes, this can work:

from __future__ import annotations
from typing import Protocol, Self, TYPE_CHECKING

if TYPE_CHECKING:
    class _MPZ_proto(Protocol):
        def __new__(cls, x: int, /) -> Self: ...
        def __int__(self) -> int: ...
        def __add__(self, other: Self, /) -> Self: ...

    class MPZ:
        def __new__(cls, x: int, /) -> Self: ...
        def __int__(self) -> int: ...
        def __add__(self, other: Self, /) -> Self: ...

    # Dummy assignments for a type checker to check that
    # MPZ, int and fmpz are defined compatibly.
    from flint import fmpz # type: ignore
    _test_MPZ1: type[_MPZ_proto] = MPZ
    _test_MPZ2: type[_MPZ_proto] = int
    _test_MPZ3: type[_MPZ_proto] = fmpz
    del _test_MPZ1, _test_MPZ2, _test_MPZ3
else:
    try:
        from flint import fmpz # type: ignore
        MPZ = fmpz
    except ImportError:
        MPZ = int

def f(x: int) -> int:
    x_mpz = MPZ(x)
    result_mpz = g(x_mpz)
    return int(result_mpz)

def g(x: MPZ) -> MPZ:
    return x + MPZ(2)

It isn’t possible for subsequent code to be typed with _MPZ_proto because then it couldn’t use MPZ(1). This checks mostly that the method definitions for MPZ are compatible with those of int but would still allow MPZ to have methods/overloads that int does not have.

9

How about using NewType?

try:
    from flint import fmpz
    MPZ = fmpz  # might need to wrap this as well
except ImportError:
    MPZ = NewType("MPZ", int)

Won’t allow you to cast other types (MPZ("123")), and will be assignable to int, but that might be fine to you

Edit: I now noticed you don’t want it to be assignable to int, so it’s not going to work