Allow type checkers to infer trait types

[vidartf]

We are currently encoding typing information for traits without integrating with the Python static typing system. This change is for allowing traitlets users to get static typing "for free" by using our trait types.

Ideally, the typing system would allow for default types, such that TraitType == TraitType[t.Any, t.Any], and TraitType[U] == TraitType[U, U], but for now we have to be verbose for this to work.

This is currently in a draft state, as some of the typings are still not accurately reflecting all the information we have. Potential improvements:

• The Union type should correctly surface its types as the typings union of those of its components
• Make the Optional part depend on allow_none (is it possible?).
• Improve types of class based traits
• Can we make the coercing variants not need to re-inherit TraitType ?
• Make Enum and UseEnum types better
• etc.

cc @maartenbreddels

[blink1073]

Nice! You may want to play with the Application object in your mypy tests, to see if we can remove the explicit typings there.

[maartenbreddels]

I've been looking into if we can make a distinction between the instance and class access:

class IntField:
    @typing.overload
    def __get__(self, instance: None=None, owner:object=...) -> "IntField":
        ...

    @typing.overload
    def __get__(self, instance: object, owner:object=...) -> int:
        ...

    def __get__(self, instance:typing.Any=None, owner:object=None) -> typing.Union["IntField", int]:
        if instance is None:
            return self
        print(instance, owner)
        return 42

    def __set__(self, instance:object, value: int) -> None:
        print("setting", value)

    def lala(self):
        print("lala")

    def on_value(self, handler: typing.Callable[[int], None]) -> None:
        pass

class Lala:
    x = IntField()

Lala.x.lala()
lala = Lala()
print(lala.x)
# lala.x = "dsa"  # gives type error :+1+
lala.x = 42

This seems to work with traitlets as well:

diff --git a/traitlets/traitlets.py b/traitlets/traitlets.py
index 0d2649e..2259289 100644
--- a/traitlets/traitlets.py
+++ b/traitlets/traitlets.py
@@ -677,7 +677,15 @@ class TraitType(BaseDescriptor, t.Generic[G, S]):
         else:
             return value  # type: ignore
 
-    def __get__(self, obj: "HasTraits", cls: t.Any = None) -> t.Optional[G]:
+    @t.overload
+    def __get__(self, obj: None=None, cls:t.Any=...) -> "TraitType[G, S]":
+        ...
+
+    @t.overload
+    def __get__(self, obj: "HasTraits", cls:t.Any=...) -> G:
+        ...
+
+    def __get__(self, obj: t.Any=None, cls: t.Any=None) -> t.Union[G, "TraitType[G, S]"]:
         """Get the value of the trait by self.name for the instance.
 
         Default values are instantiated when :meth:`HasTraits.__new__`
@@ -688,7 +696,7 @@ class TraitType(BaseDescriptor, t.Generic[G, S]):
         if obj is None:
             return self
         else:
-            return self.get(obj, cls)
+            return t.cast(G, self.get(obj, cls))
 
     def set(self, obj, value):
         new_value = self._validate(obj, value)

[vidartf]

@maartenbreddels I thought mypy already did that for class access (it isn't using __get__ type for class access, is it?).

[rmorshea]

@vidartf, it does not. MyPy checks for polymorphic behavior on __get__ to determine class vs instance access return type. Here's a sample demonstrating this.

[maartenbreddels]

it isn't using __get__ type for class access, is it?

it is, therefore the if obj is None: check

Thanks @rmorshea for the gist. Do you have an idea how to approach the allow_none -> Optional issue? Would that be possible with overload or should dataclass_transform be used for that?

[rmorshea]

This appears to work as expected. You can overload the __new__ method (just while type checking) to narrow the type. Unfortunately I don't see a clear way to avoid copy-pasting this to each TraitType because there's no way to generalize and add args to type variables (e.g. def __new__(cls: type[C], kind: T) -> C[T]: ...). Looks like there's a open issue for higher-kinded type variables.

[vidartf]

Thanks for the input. I thought I had checked it, but I must have dreamed it. Now added a test for the typing implicitly via using .tag()

[maartenbreddels]

I think with the typing_extension package we can do this:

Suggested change

	def tag(self, **metadata) -> "TraitType[G, S]":
	def tag(self, **metadata) -> typing_extension.Self":

[blink1073]

The failures are due to davidfritzsche/pytest-mypy-testing#35

[maartenbreddels]

This appears to work as expected. You can overload the __new__ method

Whow, thanks, you never stop learning with typing tricks! I think it's acceptable, the number of traits types isn't huge, and it will give us really good type checks with pylance and mypy!

[blink1073]

We can pin to pytest<7.2 for now to unblock this PR.

[blink1073]

I opened vidartf#1 exploring the idea @rmorshea had, with a bit of tweaking to get it to work with other arguments like help

[blink1073]

It turns out we will still need the explicit types on the Application object to allow subclasses to set name = "foo" for example, which is not a huge problem.