[Typing-sig] Generic `typing.ForwardRef` to support generic recursive types

21 Oct 2022

      Originally posted at [Python Discussion: Generic typing.ForwardRef to support generic recursive types](https://discuss.python.org/t/generic-typing-forwardref-to-support-generic-re...).

TL;DR: Generic recursive type annotations require `typing.ForwardRef` to support typevar substitution, like `Union`:

```python
class _GenericAlias(_Final, _root=True):
    @_tp_cache
    def __getitem__(self, params):
        if self.__origin__ in (Generic, Protocol):
            # Can't subscript Generic[...] or Protocol[...].
            raise TypeError(f"Cannot subscript already-subscripted {self}")
        if not isinstance(params, tuple):
            params = (params,)
        msg = "Parameters to generic types must be types."
        params = tuple(_type_check(p, msg) for p in params)
        _check_generic(self, params)
        return _subs_tvars(self, self.__parameters__, params)
```

Maybe we can add `__paramters__` attributes to `ForwardRef` and let it can support `_subs_tvars` in `_GenericAlias.__getitem__` (e.g. a `Union[...]`).

------

## Original Post

Since v0.981 `mypy` supports recursive types and will be enabled by default since v0.990. E.g.:

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

Recursive types need to use `ForwardRef`s at the right-hand side to reference the type alias before assignment. The ref `ForwardRef('JSON')` is not generic parameterized.

I'm trying to annotate an arbitrarily nested container with a specific element type or parameterized by a typevar. For example:

```python
T = TypeVar('T')

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

since the typevar `T` is inside a string `'PyTree[T]'`, which will be converted to a `ForwardRef`. However, it will be never evaluated because it is not assigned to a variable (`PyTree[T]` is not a valid identifier either). So I get:

```python
...
...
...
TreeOfInts = PyTree[int]
TreeOfInts
typing.Union[int, typing.Tuple[ForwardRef('PyTree[T]'), ...], typing.List[ForwardRef('PyTree[T]')], typing.Dict[typing.Any, ForwardRef('PyTree[T]')]]

`Union` only substitutes the first `T` to `int` which not in `ForwardRef`s, while the remaining `ForwardRef`s are remained as `ForwardRef('PyTree[T]')`.

As an alternative, I can specify the `int` type as:

```python
>>> TreeOfInts = Union[int, Tuple['TreeOfInts', ...], List['TreeOfInts'], Dict[Any, 'TreeOfInts']]
>>> TreeOfInts
typing.Union[int, typing.Tuple[ForwardRef('TreeOfInts'), ...], typing.List[ForwardRef('TreeOfInts')], typing.Dict[typing.Any, ForwardRef('TreeOfInts')]]

But this approach is not considered because it is impossible to cover all element types for function annotations. E.g:

```python
def tree_leaves(
    tree: PyTree[T],
    is_leaf: Optional[Callable[[T], bool]] = None,
    *,
    none_is_leaf: bool = False,
) -> List[T]: ...
```

Is it possible to make `ForwardRef` generic? Or is there any other way to define a generic recursive type (for function annotations)?

References:

* [Support recursive types python/mypy#731](https://github.com/python/mypy/issues/731)
* [Recursive Generic Type Support python/mypy#13693](https://github.com/python/mypy/issues/13693)
* [Type checking for pytrees google/jax#3340](https://github.com/google/jax/issues/3340)
* [Better Type Annotation Support for PyTrees metaopt/optree#6](https://github.com/metaopt/optree/issues/6)

[Typing-sig] Generic `typing.ForwardRef` to support generic recursive types

Xuehai Pan