Implicit TypeVar expansion of a Generic via another generic method

[chadrik]

I really struggled to come up with a title for this, so I apologize for the ambiguity. I also considered "Spooky TypeVar action at a distance".

I'm investigating strategies for static type checking of apache beam, which is a streaming pipeline API.

Basically, I have a generic Transform with an In and Out type, and another generic Collection that wants to apply its type to the input type of Transform in order to produce a new Collection with the Transform's output type. These operations are repeated in a chain to form a pipeline:

Collection[None] -> Transform[None, A] -> Collection[A] -> Transform[A, B] -> Collection[B]

The catch is that a Transform's output type can be related to its input type via a TypeVar, and a Collection contains the concrete type of that input TypeVar. So when we call Collection.apply(transform) we want to resolve that chain of TypeVar dependencies to produce an expanded/concrete output type:

from typing import *

T = TypeVar('T')
InT = TypeVar('InT')
OutT = TypeVar('OutT')


class Transform(Generic[InT, OutT]):
    "Takes an input Collection and produces an output Collection"

    def __init__(self, fn: Callable[[InT], OutT]):
        self.fn = fn

    def call(self, arg: InT) -> OutT:
        return self.fn(arg)


def get_op(f: Callable[[InT], OutT]) -> Transform[InT, OutT]:
    "Get a Transform from a callable"
    return Transform(f)


class Collection(Generic[T]):
    "Collection of elements"

    def apply(self, op: Transform[T, OutT]) -> 'Collection[OutT]':
        """
        Apply the this Collection to a Transform and produce a new 
        output Collection
        """


# -- test:

def make_string(x: None) -> str:
    return 'foo'


def make_ones(x: T) -> Tuple[T, int]:
    return (x, 1)


c1: Collection = Collection()

str_op = get_op(make_string)
reveal_type(str_op)  # revealed: Transform[None, builtins.str*]

c2 = p1.apply(str_op)
reveal_type(c2)  # revealed: Collection[builtins.str*]

tuple_op = get_op(make_ones)
reveal_type(tuple_op)  # revealed: Transform[T`-1, Tuple[T`-1, builtins.int]]

c3 = c2.apply(tuple_op)
# Desired type is Collection[Tuple[builtins.str, builtins.int]]
reveal_type(c3)  # revealed: Collection[Tuple[T`-1, builtins.int]]

The crux of the issue is here:

class Collection(Generic[T]):
    "Collection of elements"

    def apply(self, op: Transform[T, OutT]) -> 'Collection[OutT]':
        """
        Apply the this Collection to a Transform and produce a new 
        output Collection
        """

My naive desire was that my Collection[builtins.str*] would apply its str type to T-1 of Transform[T-1, Tuple[T-1, builtins.int]] to produce Collection[Tuple[builtins.str, builtins.int]].

Instead I get:

error: Argument 1 to "apply" of "Collection" has incompatible type "Transform[T, Tuple[T, int]]"; expected "Transform[str, Tuple[T, int]]"

Is there any future universe where this is possible in mypy, or is it just too ambiguous?

Is it possible to write a mypy plugin to produce the desired result for Collection.apply?

• edit: clarified relationship between Transform and Collection TypeVars

[ilevkivskyi]

Actually the fact that that T leaks its scope is a bug caused by #1317. Whether this pattern will be allowed depends on how we fix this bug. Most likely it will not be supported at least at first (otherwise this would require existential types IIUC).

I think it should possible to write a plugin however. You can try playing method hooks.

[ilevkivskyi]

I am leaving this open because it is an important variation of #1317 where return type is not a callable (see also #5738 for a similar example where a type context appears).

[chadrik]

Thanks for the info. I'll look into the plugin option soon.

[chadrik]

I started looking into how to substitute a concrete type for T into Transform[T, OutT]. I didn't find any existing plugins that are doing anything similar. contextmanager_callback, which I was hoping would be close, is much more simplistic than what I need.

I found two leads:

• mypy.typeanal.expand_type_alias(): lives in typeanal and newtypeanal, which seems off limits to plugins...
• mypy.typevars.fill_typevars(): sounds about right, but example usage of this function is not very illuminating. It doesn't seem to be used in the context of modifying an Instance, but rather for bringing an Instance into being

A nudge in the right direction would be greatly appreciated!

[ilevkivskyi]

I think we don't have many dedicated helpers for substituting some type arguments with types. You can try playing with mypy.expandtype.expand_type() (see how it is used in mypy.applytype.apply_generic_arguments()), or just write your own helper.

which seems off limits to plugins...

Technically, you can import any module you want (just be aware of import cycles).

[chadrik]

@ilevkivskyi I just realized that I forgot to thank you for your advice! (I think I typed something up but never hit the 'comment' button). mypy.expandtype.expand_type() was exactly what I was looking for.

I have written a plugin that successfully fills in the missing TypeVar using expand_type(). So now instead of Collection[Tuple[T-1, builtins.int]], I get the expected value Collection[Tuple[builtins.str, builtins.int]]`

I've hit another problem, and I'd like to know if there's another utility along the lines of expand_type() that will help.

There are scenarios where the Transform.apply() method -- and thus my plugin -- receives as an arg Transform[<nothing>, <nothing>], which I can't directly expand. I was able to restore the original type args like this:

def pvalue_or_callback(ctx: MethodContext) -> Type:
    ...
    xform_type = ctx.arg_types[0][0]
    xform_type_inhabited = xform_type.copy_modified(
        args=xform_type.type.bases[0].args)
    ...

That gets xform_type_inhabited back to this:

Transform[Tuple[K_`1, V_`2], Tuple[K_`1, Iterable[V_`2]]]

Now I want to apply Tuple[str, int] to those args, such that I end up with this:

Transform[Tuple[str, int], Tuple[str, Iterable[int]]]

Is there an easy way to do this?

Thanks again!

[ilevkivskyi]

Is there an easy way to do this?

Hm, I don't think we have any dedicated helpers for such cases.