Improve generics handling (#886)

ipython · Oct 30, 2023 · dacbf9b · dacbf9b
1 parent 4f0e68b
commit dacbf9b
Show file tree

Hide file tree

Showing 4 changed files with 174 additions and 98 deletions.
diff --git a/tests/test_typing.py b/tests/test_typing.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import logging
 import typing as t
+from abc import ABC
 
 import pytest
 
@@ -9,6 +11,7 @@
     Bool,
     CInt,
     Dict,
+    Enum,
     HasTraits,
     Instance,
     Int,
@@ -97,12 +100,14 @@ class T(HasTraits):
         ).tag(config=True)
 
     t = T()
-    reveal_type(List("foo"))  # R: traitlets.traitlets.List
-    reveal_type(List("").tag(sync=True))  # R: traitlets.traitlets.List
-    reveal_type(List(None, allow_none=True))  # R: traitlets.traitlets.List
-    reveal_type(List(None, allow_none=True).tag(sync=True))  # R: traitlets.traitlets.List
-    reveal_type(T.latex_command)  # R: traitlets.traitlets.List
-    reveal_type(t.latex_command)  # R: builtins.list[Any]
+    reveal_type(List(["foo"]))  # R: traitlets.traitlets.List[builtins.str]
+    reveal_type(List([""]).tag(sync=True))  # R: traitlets.traitlets.List[builtins.str]
+    reveal_type(List(None, allow_none=True))  # R: traitlets.traitlets.List[<nothing>]
+    reveal_type(
+        List(None, allow_none=True).tag(sync=True)  # R: traitlets.traitlets.List[<nothing>]
+    )
+    reveal_type(T.latex_command)  # R: traitlets.traitlets.List[builtins.str]
+    reveal_type(t.latex_command)  # R: builtins.list[builtins.str]
 
 
 @pytest.mark.mypy_testing
@@ -111,19 +116,25 @@ class T(HasTraits):
         foo = Dict({}, help="Shell command used to compile latex.").tag(config=True)
 
     t = T()
-    reveal_type(Dict("foo"))  # R: traitlets.traitlets.Dict
-    reveal_type(Dict("").tag(sync=True))  # R: traitlets.traitlets.Dict
-    reveal_type(Dict(None, allow_none=True))  # R: traitlets.traitlets.Dict
-    reveal_type(Dict(None, allow_none=True).tag(sync=True))  # R: traitlets.traitlets.Dict
-    reveal_type(T.foo)  # R: traitlets.traitlets.Dict
-    reveal_type(t.foo)  # R: builtins.dict[Any, Any]
+    reveal_type(Dict(None, allow_none=True))  # R: traitlets.traitlets.Dict[builtins.str, Any]
+    reveal_type(
+        Dict(None, allow_none=True).tag(sync=True)  # R: traitlets.traitlets.Dict[builtins.str, Any]
+    )
+    reveal_type(T.foo)  # R: traitlets.traitlets.Dict[builtins.str, Any]
+    reveal_type(t.foo)  # R: builtins.dict[builtins.str, Any]
 
 
 @pytest.mark.mypy_testing
 def mypy_type_typing() -> None:
     class KernelSpec:
         item = Unicode("foo")
 
+    class KernelSpecSubclass(KernelSpec):
+        other = Unicode("bar")
+
+    class GatewayTokenRenewerBase(ABC):
+        item = Unicode("foo")
+
     class KernelSpecManager(HasTraits):
         """A manager for kernel specs."""
 
@@ -136,12 +147,30 @@ class KernelSpecManager(HasTraits):
         )
         other_class = Type("foo.bar.baz")
 
+        other_kernel_spec_class = Type(
+            default_value=KernelSpecSubclass,
+            klass=KernelSpec,
+            config=True,
+        )
+
+        gateway_token_renewer_class = Type(
+            klass=GatewayTokenRenewerBase,
+            config=True,
+            help="""The class to use for Gateway token renewal. (JUPYTER_GATEWAY_TOKEN_RENEWER_CLASS env var)""",
+        )
+
     t = KernelSpecManager()
-    reveal_type(t.kernel_spec_class)  # R: def () -> tests.test_typing.KernelSpec@124
-    reveal_type(t.kernel_spec_class())  # R: tests.test_typing.KernelSpec@124
+    reveal_type(t.kernel_spec_class)  # R: def () -> tests.test_typing.KernelSpec@129
+    reveal_type(t.kernel_spec_class())  # R: tests.test_typing.KernelSpec@129
     reveal_type(t.kernel_spec_class().item)  # R: builtins.str
     reveal_type(t.other_class)  # R: builtins.type
     reveal_type(t.other_class())  # R: Any
+    reveal_type(t.other_kernel_spec_class)  # R: def () -> tests.test_typing.KernelSpec@129
+    reveal_type(t.other_kernel_spec_class())  # R: tests.test_typing.KernelSpec@129
+    reveal_type(
+        t.gateway_token_renewer_class  # R: def () -> tests.test_typing.GatewayTokenRenewerBase@135
+    )
+    reveal_type(t.gateway_token_renewer_class())  # R: tests.test_typing.GatewayTokenRenewerBase@135
 
 
 @pytest.mark.mypy_testing
@@ -181,6 +210,42 @@ class T(HasTraits):
     reveal_type(t.export_format)  # R: builtins.str
 
 
+@pytest.mark.mypy_testing
+def mypy_enum_typing() -> None:
+    class T(HasTraits):
+        log_level = Enum(
+            (0, 10, 20, 30, 40, 50),
+            default_value=logging.WARN,
+            help="Set the log level by value or name.",
+        ).tag(config=True)
+
+    t = T()
+    reveal_type(
+        Enum(  # R: traitlets.traitlets.Enum[builtins.str]
+            ("foo",)
+        )
+    )
+    reveal_type(
+        Enum(  # R: traitlets.traitlets.Enum[builtins.str]
+            [""]
+        ).tag(sync=True)
+    )
+    reveal_type(
+        Enum(  # R: traitlets.traitlets.Enum[None]
+            None, allow_none=True
+        )
+    )
+    reveal_type(
+        Enum(  # R: traitlets.traitlets.Enum[None]
+            None, allow_none=True
+        ).tag(sync=True)
+    )
+    reveal_type(
+        T.log_level  # R: traitlets.traitlets.Enum[builtins.int]
+    )
+    reveal_type(t.log_level)  # R: builtins.int
+
+
 @pytest.mark.mypy_testing
 def mypy_set_typing() -> None:
     class T(HasTraits):

diff --git a/traitlets/config/application.py b/traitlets/config/application.py
@@ -199,7 +199,7 @@ def _classes_inc_parents(
     version: str | Unicode[str, str | bytes] = Unicode("0.0")
 
     # the argv used to initialize the application
-    argv = List()
+    argv: list[str] | List[str] = List()
 
     # Whether failing to load config files should prevent startup
     raise_config_file_errors = Bool(TRAITLETS_APPLICATION_RAISE_CONFIG_FILE_ERROR)
@@ -241,7 +241,7 @@ def get_default_logging_config(self) -> StrDict:
                 "console": {
                     "class": "logging.StreamHandler",
                     "formatter": "console",
-                    "level": logging.getLevelName(self.log_level),
+                    "level": logging.getLevelName(self.log_level),  # type:ignore[arg-type]
                     "stream": "ext://sys.stderr",
                 },
             },
@@ -278,7 +278,7 @@ def _observe_logging_change(self, change: Bunch) -> None:
         # convert log level strings to ints
         log_level = self.log_level
         if isinstance(log_level, str):
-            self.log_level = getattr(logging, log_level)
+            self.log_level = t.cast(int, getattr(logging, log_level))
         self._configure_logging()
 
     @observe("log", type="default")
@@ -400,7 +400,7 @@ def _log_default(self) -> AnyLogger:
     # this must be a dict of two-tuples,
     # the first element being the application class/import string
     # and the second being the help string for the subcommand
-    subcommands: dict[str, t.Any] | Dict = Dict()
+    subcommands: dict[str, t.Any] | Dict[str, t.Any] = Dict()
     # parse_command_line will initialize a subapp, if requested
     subapp = Instance("traitlets.config.application.Application", allow_none=True)
 
@@ -418,7 +418,7 @@ def _log_default(self) -> AnyLogger:
         """,
     )
 
-    _loaded_config_files = List()
+    _loaded_config_files: List[str] = List()
 
     show_config = Bool(
         help="Instead of starting the Application, dump configuration to stdout"

diff --git a/traitlets/config/loader.py b/traitlets/config/loader.py
@@ -98,9 +98,9 @@ class LazyConfigValue(HasTraits):
     _value = None
 
     # list methods
-    _extend: List = List()
-    _prepend: List = List()
-    _inserts: List = List()
+    _extend: List[t.Any] = List()
+    _prepend: List[t.Any] = List()
+    _inserts: List[t.Any] = List()
 
     def append(self, obj: t.Any) -> None:
         """Append an item to a List"""