pytorch/tnt

Use PEP692 for AutoUnit kwargs typing

alanhdu opened this issue ยท 1 comments

๐Ÿš€ The feature

Would it be possible to define TypedDicts corresponding to the kwargs for AutoUnit.

Motivation, pitch

Right now, when you implementers of AutoUnit are need to subclass AutoUnit and then call super().__init__ -- because AutoUnit itself takes many different flags, it's fairly normal to do:

class MySubclass(AutoUnit):
    def __init__(self, my_arg: whatever, ..., **kwargs):
        super().__init__(**kwargs)
        # do stuff with `my_arg`

Unfortunately, this loses typehints on the **kwargs. PEP 692 allows you to use Unpack with a TypedDict -- assuming that we defined AutoUnitKwargs(TypedDict) appropriately, this would allow implementers to do:

class MySubclass(AutoUnit):
    def __init__(self, my_arg: whatever, ..., **kwargs: Unpack[AutoUnitKwargs])

which would preserve the type hints. This would be much more convenient than having to replicate AutoUnit's type signature and manually forwarding them to super().__init__.

Alternatives

The main alternative is to just do nothing, and either live with untyped kwargs or ask users to duplicate the typehints for every subclass.

Additional context

No response

Hey Alan, thanks for using TorchTNT and for this suggestion. I agree that it will help users write safer code.

Out of curiosity, which type checker are you using? because it looks like Pyre doesn't support preserving the type hints in Unpack with TypedDict just yet.

Once the support will be added in Pyre, we can definitely revisit to add this.