Generic `typing.ForwardRef` to support generic recursive types

Ideas
,
1

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

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

Recursive types need to use ForwardRefs 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:

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:

>>> 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 ForwardRefs, while the remaining ForwardRefs are remained as ForwardRef('PyTree[T]').

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

>>> 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:

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:

2 Likes
2

This looks like an interesting feature proposal.

A better venue for discussing it is probably typing-sig@python.org (join the list at Mailman 3 Info | typing-sig@python.org - python.org).

2 Likes