google-deepmind / tree

tree is a library for working with nested data structures
https://tree.readthedocs.io
Apache License 2.0
933 stars 57 forks source link

Add type hints to public APIs #84

Open ikamensh opened 2 years ago

ikamensh commented 2 years ago

Hello,

I'd like to propose a feature - to add type hints to public functions.

The most basic version of this would be to type hint return types in terms of TypeVar of input types, like:

S = TypeVar("S")  # type of a structure
def map_structure(func, *structures: S, **kwargs) -> S:
  ...

This could be taken further by making structure generic over it's item types, allowing to type hint functions like flatten:

Item = TypeVar("Item")
Node = Union["Structure[Item]", Item]

Structure = Union[
    Node,
    Sequence['Node'],
    Mapping[Any, 'Node'],
    'AnyNamedTuple',
]

def flatten(structure: Structure[Item]) -> list[Item]:
  ...

Overall, this should help users avoid preventable bugs and have more readable code. Thanks for considering :)

superbobry commented 1 year ago

Thank for the suggestion Ilya. Please feel free to send a PR for this.

My gut feeling is that it could be tricky to come up with useful types for a lot of the APIs, because they are highly polymorphic, but some type information is still definitely better than none at all.