Fragile code and lots of trial and error

[Massimiliano Tomassoli (mtomassoli)]

Maybe I'm asking too much of Pyright, but even when I find a mathematically correct way to get the result I want, I have then to try out lots of equivalent formulations to make it work.

Here's an example. Suppose we have unlifted and lifted types. An unlifted type is a built-in or pre-existing type, whereas the associated lifted type is a type we created and thus control.

For instance, list[T] is an unlifted type, and my_list[T] is the corresponding lifted type.

Let's say the function that does the lifting is called lift: it takes an unlifted object and returns the corresponding lifted object.

A simple solution to write lift would be to create an overload for each "unlifted_type -> lifted_type" correspondence, but that's inefficient. If lift has 1k overloads, 1k uses take minutes on my PC.

I did manage to bring the time down to 2-3 secs by building an n-ary tree of types and then doing an n-ary search.

The code is a little involved. Here's an example for 1k elements:

# 3k LOC of type definitions omitted

T1 = TypeVar('T1', contravariant=True)
T2 = TypeVar('T2')
T3 = TypeVar('T3')

class C(Generic[T1, T2, T3]): ...

def __sel9[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel8[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel7[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel6[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel5[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel4[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel3[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel2[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel1[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...
def __sel0[U1, G1, G2](g1: G1, g2: G2, x: C[G1, U1, Any] | C[G2, Any, U1]) -> U1: ...

def sel[OBJ,
        C8_1, C8_2, C7_1, C7_2, C6_1, C6_2, C5_1, C5_2, C4_1, C4_2, C3_1, C3_2, C2_1, C2_2, C1_1, C1_2, C0_1, C0_2, C0,
        T8_1, T8_2, T7_1, T7_2, T6_1, T6_2, T5_1, T5_2, T4_1, T4_2, T3_1, T3_2, T2_1, T2_2, T1_1, T1_2](
    obj: OBJ, *,
    __s9: Callable[[Cla9_1, Cla9_2, C[OBJ, Tup9_1, Tup9_2]],
                   tuple[C8_1, C8_2, T8_1, T8_2]] = __sel9,    # type: ignore
    __s8: Callable[[C8_1, C8_2, C[OBJ, T8_1, T8_2]],
                   tuple[C7_1, C7_2, T7_1, T7_2]] = __sel8,    # type: ignore
    __s7: Callable[[C7_1, C7_2, C[OBJ, T7_1, T7_2]],
                   tuple[C6_1, C6_2, T6_1, T6_2]] = __sel7,    # type: ignore
    __s6: Callable[[C6_1, C6_2, C[OBJ, T6_1, T6_2]],
                   tuple[C5_1, C5_2, T5_1, T5_2]] = __sel6,    # type: ignore
    __s5: Callable[[C5_1, C5_2, C[OBJ, T5_1, T5_2]],
                   tuple[C4_1, C4_2, T4_1, T4_2]] = __sel5,    # type: ignore
    __s4: Callable[[C4_1, C4_2, C[OBJ, T4_1, T4_2]],
                   tuple[C3_1, C3_2, T3_1, T3_2]] = __sel4,    # type: ignore
    __s3: Callable[[C3_1, C3_2, C[OBJ, T3_1, T3_2]],
                   tuple[C2_1, C2_2, T2_1, T2_2]] = __sel3,    # type: ignore
    __s2: Callable[[C2_1, C2_2, C[OBJ, T2_1, T2_2]],
                   tuple[C1_1, C1_2, T1_1, T1_2]] = __sel2,    # type: ignore
    __s1: Callable[[C1_1, C1_2, C[OBJ, T1_1, T1_2]],
                   tuple[C0_1, C0_2]] = __sel1,    # type: ignore
    __s0: Callable[[C0_1, C0_2, C[OBJ, C0_1, C0_2]],
                   C0] = __sel0,
    ) -> C0: ...

This code is generated programmatically, of course :)

I've experimented with lots of other ideas, but this is one of the most performant.

I then realized that I had to import all the unlifted and lifted types to make this work, so I gave up on this approach.

The new idea is to create small custom lift right when they're needed.

For instance, if a file uses 5 unlifted types, then one can just add support for them this way:

lift = Lift[A1[Any] | A2[Any] | A3[Any] | A4[Any] | A5[Any]]()

A1, ..., A5 are the lifted types, and the Anys can be omitted.

One may then lift a supported unlifted type this way:

a1 = lift(Unlifted1[int]())
assert_type(a1, A1[int])

Again, all this is necessary because we have no control over Unlifted1 so how can we find the corresponding A1? We have to search for it, of course.

Here's the full code, including the tests:

from __future__ import annotations
from typing import (
    Protocol, TypeVar, Generic, Any, Callable, TYPE_CHECKING, assert_type
)

class LiftedTypeP(Protocol):
    _mut_unlifted_type: ...
    _mut_lifted_type: ...
    _mut_lift: ...

LI = TypeVar('LI')
UT = TypeVar('UT')
LT = TypeVar('LT')

# NOTE: `UT` MUST be INVARIANT or unlifted types may be hidden by unlifted
#   supertypes. (Pyright BUG?: this doesn't work when unlifted types have type
#   params.)
class _LiftedTypeP(Generic[UT, LT, LI], Protocol):   # pyright: ignore[reportInvalidTypeVarUse]
    _mut_unlifted_type: UT
    _mut_lifted_type: LT
    _mut_lift: LI

# NOTE: `UT` MUST be INVARIANT.
class _Shell(Generic[UT]): ...

# NOTE: `UT` is protected by `_Shell` so that `UT` can't be the union of two or
#   more unlifted types.
#   Without this, an `unlifted` of type `Unlifted1` might, for example, give
#       UT = Unlifted1 | Unlifted4;     LT = Lifted4
#   Thanks to `_Shell`, the only lecit solution is
#       UT = Unlifted1;                 LT = Lifted1
def _f1[UT, LT](unlifted: _Shell[UT], x: _LiftedTypeP[UT, LT, Any] | Any
                ) -> LT: ...
def _f2[LI](x: _LiftedTypeP[Any, Any, LI]) -> LI: ...
def _f3[UT, LT](ut: UT, f: Callable[[UT], LT]) -> LT: ...

class Lift[LTS]:
    # Pyright BUG?: I had to split `_f` into 3 functions and use them in the
    # "right" order.
    def __call__[UT, LT, LT2, LT3](
        self, unlifted: UT, *,
        __f3: Callable[[UT, LT2], LT3] = _f3,
        __f2: Callable[[LT], LT2] = _f2,
        __f: Callable[[_Shell[UT], LTS], LT] = _f1,
        ) -> LT3: ...

#----------------------------------- TESTS ------------------------------------

if TYPE_CHECKING:
    class Unlifted1[T]: ...
    class Unlifted2[T]: ...
    class Unlifted3[T]: ...
    class Unlifted4[T]: ...
    class Unlifted5[T]: ...
    
    class A1[T](LiftedTypeP):
        _mut_unlifted_type: Unlifted1[T]
        _mut_lifted_type: A1[T]
        def _mut_lift[U](self, unl: Unlifted1[U]) -> A1[U]: ...

    class A2[T](LiftedTypeP):
        _mut_unlifted_type: Unlifted2[T]
        _mut_lifted_type: A2[T]
        def _mut_lift[U](self, unl: Unlifted2[U]) -> A2[U]: ...

    class A3[T](LiftedTypeP):
        _mut_unlifted_type: Unlifted3[T]
        _mut_lifted_type: A3[T]
        def _mut_lift[U](self, unl: Unlifted3[U]) -> A3[U]: ...

    class A4[T](LiftedTypeP):
        _mut_unlifted_type: Unlifted4[T]
        _mut_lifted_type: A4[T]
        def _mut_lift[U](self, unl: Unlifted4[U]) -> A4[U]: ...

    class A5[T](LiftedTypeP):
        _mut_unlifted_type: Unlifted5[T]
        _mut_lifted_type: A5[T]
        def _mut_lift[U](self, unl: Unlifted5[U]) -> A5[U]: ...

    lift = Lift[A1[Any] | A2[Any] | A3[Any] | A4[Any] | A5[Any]]()

    assert_type(lift(Unlifted1[int]()), A1[int])
    assert_type(lift(Unlifted2[str]()), A2[str])
    assert_type(lift(Unlifted3[list[int]]()), A3[list[int]])
    assert_type(lift(Unlifted4[set[Any]]()), A4[set[Any]])
    assert_type(lift(Unlifted5[dict[int, str]]()), A5[dict[int, str]])

My problem with this is that the method __call__ right before the tests is extremely fragile. Getting it to work was extremely painful, and every little change breaks the code. For instance, swapping 2 __f? is enough to break it.

The versions with a single f didn't work, the ones with 2 fs worked but gave weird error messages, and, finally, the current version with 3 fs works. But it feels like I'm exploiting implementation details.

Am I doing something wrong or are these limitations of Pyright?

[Eric Traut (erictraut)]

You're abusing the Python static type system here. It's unclear to me what you're trying to do, but I can say with high confidence that this approach won't work well. This is not a use case that pyright is designed to handle. I think it's safe to say the same about other Python static type checkers like mypy. I recommend looking for an alternate approach to whatever problem you're trying to solve.