Recursive Generic Type Support

[XuehaiPan]

Bug Report

In:

• Support recursive types #731
• Enable recursive type aliases behind a flag #13297 (not released yet in v0.971)

We now support recursive type hints, such as:

JSON = Union[Dict[str, 'JSON'], List['JSON'], str, int, float, bool, None]

But in the example in #13297 (comment), which is generic and has typevars:

Nested = Sequence[Union[T, Nested[T]]]

def flatten(seq: Nested[T]) -> List[T]:
    flat: List[T] = []
    for item in seq:
        if isinstance(item, Sequence):
            res.extend(flatten(item))
        else:
            res.append(item)
    return flat

reveal_type(flatten([1, [2, [3]]]))  # N: Revealed type is "builtins.list[builtins.int]"

The code is Python syntax illegal that the name Nested in Nested[T] is referenced before assignment.

My own use case: define a generic pytree type:

"""pytree_forwardref.py"""

from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TypeVar, Union, Protocol

import torch

T = TypeVar('T')

Children = Sequence[T]
_AuxData = TypeVar('_AuxData', bound=Hashable)
AuxData = Optional[_AuxData]


class CustomTreeNode(Protocol[T]):
    """The abstract base class for custom pytree nodes."""

    def tree_flatten(self) -> Tuple[Children[T], AuxData]:
        """Flattens the custom pytree node into children and auxiliary data."""

    @classmethod
    def tree_unflatten(cls, aux_data: AuxData, children: Children[T]) -> 'CustomTreeNode[T]':
        """Unflattens the children and auxiliary data into the custom pytree node."""


PyTree = Union[
    T,
    Tuple['PyTree[T]', ...],  # Tuple, NamedTuple
    List['PyTree[T]'],
    Dict[Any, 'PyTree[T]'],  # Dict, OrderedDict, DefaultDict
    CustomTreeNode['PyTree[T]'],
]


TensorTree = PyTree[torch.Tensor]
print(TensorTree)

I got NameError is use PyTree[T]. Or the typevar T is not expended when using ForwardRef ('PyTree[T]'):

$ python3 pytree_forwardref.py                      
typing.Union[torch.Tensor, typing.Tuple[ForwardRef('PyTree[T]'), ...], typing.List[ForwardRef('PyTree[T]')], typing.Dict[typing.Any, ForwardRef('PyTree[T]')], __main__.CustomTreeNode[ForwardRef('PyTree[T]')]]

To Reproduce

1. mypy does not report NameError for recursive type:

   1. Install pylint and the dev version of mypy:

   pip3 install pylint
   pip3 install git+https://github.com/python/mypy.git

   2. Create a file with content:

   """nested.py"""
   
   from typing import TypeVar, Sequence, Union
   
   T = TypeVar('T')
   
   Nested = Sequence[Union[T, Nested[T]]]
   
   NestedInt = Nested[int]
   print(NestedInt)

   3. Run mypy:

   $ mypy --enable-recursive-aliases nested.py
   Success: no issues found in 1 source file

   4. Run pylint:

   $ pylint nested.py 
   ************* Module nested
   nested.py:7:27: E0601: Using variable 'Nested' before assignment (used-before-assignment)
   
   ------------------------------------------------------------------
   Your code has been rated at 0.00/10 (previous run: 0.00/10, +0.00)
   
   $ python3 nested.py
   Traceback (most recent call last):
     File "nested.py", line 7, in <module>
       Nested = Sequence[Union[T, Nested[T]]]
   NameError: name 'Nested' is not defined

2. ForwardRef does not expand typevars:

   5. Change to use ForwardRef:

   """nested_forwardref.py"""
   
   from typing import TypeVar, Sequence, Union
   
   T = TypeVar('T')
   
   Nested = Sequence[Union[T, 'Nested[T]']]
   
   NestedInt = Nested[int]
   print(NestedInt)

   6. Run mypy, pylint and execute:

   $ mypy --enable-recursive-aliases nested_forwardref.py
   Success: no issues found in 1 source file
   
   $ pylint nested_forwardref.py
   
   ------------------------------------
   Your code has been rated at 10.00/10
   
   $ python3 python3 nested_forwardref.py
   typing.Sequence[typing.Union[int, ForwardRef('Nested[T]')]]

Expected Behavior

1. Raise errors for generic recursive types.

or

2. support generic ForwardRef (preferred, but maybe need support by Python)

Actual Behavior

No error raise when using the variable name without a ForwardRef.

Ref:

• Python Discussion: Generic typing.ForwardRef to support generic recursive types.

Your Environment

• Mypy version used: master @ 1d4395f
• Mypy command-line flags: --enable-recursive-aliases
• Mypy configuration options from mypy.ini (and other config files): None
• Python version used: 3.8.12
• Operating system and version: Ubuntu 20.04 LTS