Expression dependant generic types

[plannigan]

Consider a generic tree and a transform class (produces some result after visiting every node).

class Tree(Generic[_Child_T]):
    children: List[Union[Tree[_Child_T], _Child_T]]

class Transformer(Generic[_Child_T, _Return_T]):
    def transform(self, tree: Tree[_Child_T]) -> _Return_T:
       ...

Note that _Return_T could be anything, but it is common for it to be a Tree. As a result, it is convenient cont chain Transformers together. This would allow the code to go from Tree[T] to an intermediary Tree[U], but then finally to V

class Transformer(Generic[_Child_T, _Return_T]):
    def transform(self, tree: Tree[_Child_T]) -> _Return_T:
        ...
   def chain(self, other: Transformer[_Child_U, _Return_V]) -> Transformer[_Child_T, _Return_V]:
        ...

The catch here is that this will only work when _Return_T is Tree[_Child_U]. Is it possible to express this with the typing system?

[srittau]

What happens at runtime when _ReturnT is not Tree[_Child_U]? Will chain() raise an exception? Or would chain be implemented in a subclass ChainableTransformer that must have an appropriate transform()?

[plannigan]

The current implementation, that I'm attempting to improve the type annotations for, does use a TransformerChain class that looks like this:

class TransformerChain:
    def __init__(self, transformers):
        self.transformers = transformers

    def transform(self, tree):
        for t in self.transformers:
            tree = t.transform(tree)
        return tree

So there currently aren't any checks happening in chain() or TransformerChain.transform(). So the result would be some exception raised when transform() is eventually called. I'd like to be able to detect this type of error statically with mypy.

[sobolevn]

Will something like

def chain(self, other: Transformer[_Child_U, Tree[_Return_V]]) -> Transformer[_Child_T, Tree[_Return_V]]:
        ...

work for you?

[plannigan]

That isn't exactly right. The actual implementation is like other.transform(self.transform(tree)).

So that chain suggestion would enforce that other produces a new Tree. However, that isn't what is required here as it needs to enforce that a Tree is produced by the original Transformer. But your comment did pushed me in the correct direction. While generally, there isn't a need to annotate self for each method, but that is exactly what I want here.

[plannigan]

Thanks to @sobolevn's comment, I was able to come up with the answer I was looking for. The key is to annotate the self argument to the chain() method, despite that rarely being necessary. The following code say that for a call to chain() to be valid _Return_T == Tree[_Child_U], but does not require that all Transformers produce a Tree.

from typing import Generic, List, TypeVar, Union

_Child_T = TypeVar("_Child_T")
_Child_U = TypeVar("_Child_U")
_Return_T = TypeVar("_Return_T")
_Return_V = TypeVar("_Return_V")


class Tree(Generic[_Child_T]):
    children: "List[Union[Tree[_Child_T], _Child_T]]"


class Transformer(Generic[_Child_T, _Return_T]):
    def transform(self, tree: Tree[_Child_T]) -> _Return_T:
        pass

    def chain(
            self: "Transformer[_Child_T, Tree[_Child_U]]",
            other: "Transformer[_Child_U, _Return_V]"
    ) -> "Transformer[_Child_T, _Return_V]":
        pass


tree_str_to_int: Transformer[str, int] = Transformer()
tree_str_to_tree_int: Transformer[str, Tree[int]] = Transformer()
tree_str_to_str: Transformer[str, str] = Transformer()
tree_int_to_str: Transformer[int, str] = Transformer()

tree_str_to_int.chain(tree_str_to_int)
tree_str_to_tree_int.chain(tree_str_to_str)
tree_str_to_tree_int.chain(tree_int_to_str)

When checked with mypy the results are:

$ mypy --version
mypy 0.910
$ mypy typing_example.py 
typing_example.py:29: error: Invalid self argument "Transformer[str, int]" to attribute function "chain" with type "Callable[[Transformer[_Child_T, Tree[_Child_U]], Transformer[_Child_U, _Return_V]], Transformer[_Child_T, _Return_V]]"
typing_example.py:29: error: Argument 1 to "chain" of "Transformer" has incompatible type "Transformer[str, int]"; expected "Transformer[<nothing>, int]"
typing_example.py:30: error: Argument 1 to "chain" of "Transformer" has incompatible type "Transformer[str, str]"; expected "Transformer[int, str]"
Found 3 errors in 1 file (checked 1 source file)

• The first chain() call is not valid because the self reference's _Return_T isn't Tree. The error message could be a bit clearer, but it works. (The line also produces a second error. I guess mypy isn't 100% clear what the real issue is, so it prints multiple message.)
• While the second chain() call does have a self reference that produces a Tree, it is not valid because the other argument expects a different child type for the Tree.
• The final chain() call is valid. Tree[str] -> Tree[int] -> str