Going off the rails wrangling Python's type system

Offtopic
1

I love types… maybe too much. It’s definitely my favorite thing about Julia.

I recently made the type checker much stricter on my Python projects trying to get better clarity about my Python code. But that means repeating over and over

y = f(x)
if not isinstance(y, ExpectedType):
    raise TypeError("We shouldn't be here.")

Or sometimes I want to coerce to a type, so it’s like

y = f(x)
if not isinstance(y, ExpectedType):
    y = some_fix_func(y)
if not isinstance(y, ExpectedType):
    raise TypeError("We shouldn't be here.")

I set out to write a function that would capture this behavior and be transparent to the type checker and it is… it’s not simple to do it. With no type-level computation, no issubtype on generics in at type check time, and UnionType NOT BEING A TYPE (grrr…), there were a bunch of hurdles.

But now, if you write any Python and would like to insert nice type checks that work at check time and run time, have I got a hacky class for you!

from __future__ import annotations

from types import GenericAlias
from typing import Callable, Generic, TypeVar, overload

T = TypeVar("T")
X = TypeVar("X")
T1 = TypeVar("T1")
T2 = TypeVar("T2")
T3 = TypeVar("T3")


class _Typed(Generic[T]):
    """Construct a callable type checking function.

    Instead
    """

    def __init__(self, ts: tuple = ()) -> None:
        self.ts = ts

    @staticmethod
    def _strip_generics(t: GenericAlias | type[T1]) -> type[T1]:
        """Get origin type of generic."""
        if isinstance(t, GenericAlias):
            return t.__origin__  # pyright: ignore[reportReturnType]
        return t

    @overload
    def __getitem__(self, t: type[T1]) -> _Typed[T1]: ...
    @overload
    def __getitem__(self, t: tuple[type[T1], type[T2]]) -> _Typed[T1 | T2]: ...
    @overload
    def __getitem__(
        self,
        t: tuple[type[T1], type[T2], type[T3]],
    ) -> _Typed[T1 | T2 | T3]: ...
    def __getitem__(self, t):
        if not isinstance(t, tuple):
            t = (t,)
        ts = tuple(self._strip_generics(_t) for _t in t)
        return _Typed(ts)

    @overload
    def __call__(
        self,
        x: T | X,
        /,
        *,
        recovery: Callable[[X], T] | None = None,
        err_msg: str = "",
    ) -> T: ...
    @overload
    def __call__(
        self,
        x: X,
        /,
        *,
        recovery: Callable[[X], T] | None = None,
        err_msg: str = "",
    ) -> T: ...

    def __call__(
        self,
        x,
        /,
        *,
        recovery=None,
        err_msg="",
    ) -> T:
        """Enforce type constraints on `x` with parameters of `typed`.

        `typed[T](x)` will ensure that x is a T.

        Kwargs:
        - recovery (Optional[callable]): a function to call on `x` if it is not
            the correct type.
        - err_msg (Optional[str]): a message to prepend to the type error message.
        """
        for t in self.ts:
            if isinstance(x, t):
                return x
        if recovery is not None:
            return self(
                recovery(x),
                err_msg=err_msg,
            )
        if len(err_msg) > 0:
            err_msg += "\n"
        type_str = " | ".join(self.ts)
        err_msg += f"Expected type {type_str}, but got type {type(x)}"
        raise TypeError(err_msg)


typed = _Typed()

With this class, we can do:

y = typed[ExpectedType](f(x), recovery=some_fix_func)

And now the type checker is convinced that y is ExpectedType.

2 Likes
2

Interesting have you tried mojo too ?

3

I’m not terribly interested in learning a proprietary language modelled after a language that gives me this much grief :sweat_smile: If I need to reach for performance in a Python project, calling out to Julia works just fine for me.

Why do you ask?

4

Its just that types in mojo looks a bit more well types and it may help in this situation but I have no idea how it interopt with python types, I agree with you that going to julia is always a good way

1 Like
5

Congratulations you have reinvented Validation Decorator - Pydantic.

1 Like
6

So, I saw this, but it didn’t solve my problem because I don’t want to have to wrap every library call that could return Some | None in a validated function. Getting this working as a generic function (instead of a decorator) was my goal.

The really annoying thing is that if you just do typed(t: type[T], obj: X, recovery: Callable[[X], T) -> T, and then pass a recovery function that returns the wrong type, the type check widens the definition of T to include the return type, so you don’t get a warning. That’s why I had to split it into a class definition and a function call – that forced the type parameter to be defined by t and not recovery’s return type.