Make Set, Tuple, Dict, and UseEnum generic.

Author: Alan

tests/test_typing.py

@@ -249,16 +249,13 @@ class T(HasTraits):
 @pytest.mark.mypy_testing
 def mypy_set_typing() -> None:
     class T(HasTraits):
-        remove_cell_tags = Set(
+        remove_cell_tags: Set[str] = Set(
             Unicode(),
             default_value=[],
-            help=(
-                "Tags indicating which cells are to be removed,"
-                "matches tags in ``cell.metadata.tags``."
-            ),
+            help=("Tags indicating which cells are to be removed,matches tags in ``cell.metadata.tags``."),
         ).tag(config=True)
 
-        safe_output_keys = Set(
+        safe_output_keys: Set[t.Any] = Set(
             config=True,
             default_value={
                 "metadata",  # Not a mimetype per-se, but expected and safe.
@@ -272,13 +269,13 @@ class T(HasTraits):
         )
 
     t = T()
-    reveal_type(Set("foo"))  # R: traitlets.traitlets.Set
-    reveal_type(Set("").tag(sync=True))  # R: traitlets.traitlets.Set
-    reveal_type(Set(None, allow_none=True))  # R: traitlets.traitlets.Set
-    reveal_type(Set(None, allow_none=True).tag(sync=True))  # R: traitlets.traitlets.Set
-    reveal_type(T.remove_cell_tags)  # R: traitlets.traitlets.Set
-    reveal_type(t.remove_cell_tags)  # R: builtins.set[Any]
-    reveal_type(T.safe_output_keys)  # R: traitlets.traitlets.Set
+    reveal_type(Set({"foo"}))  # R: traitlets.traitlets.Set[builtins.str]
+    reveal_type(Set({""}).tag(sync=True))  # R: traitlets.traitlets.Set[builtins.str]
+    reveal_type(Set(None, allow_none=True))  # R: traitlets.traitlets.Set[Never]
+    reveal_type(Set(None, allow_none=True).tag(sync=True))  # R: traitlets.traitlets.Set[Never]
+    reveal_type(T.remove_cell_tags)  # R: traitlets.traitlets.Set[builtins.str]
+    reveal_type(t.remove_cell_tags)  # R: builtins.set[builtins.str]
+    reveal_type(T.safe_output_keys)  # R: traitlets.traitlets.Set[Any]
     reveal_type(t.safe_output_keys)  # R: builtins.set[Any]
 
 

traitlets/traitlets.py

@@ -495,16 +495,15 @@ def instance_init(self, obj: t.Any) -> None:
 G = TypeVar("G")
 S = TypeVar("S")
 T = TypeVar("T")
-
+K = TypeVar("K", default=str)
+V = TypeVar("V", default=t.Any)
 
 # Self from typing extension doesn't work well with mypy https://github.com/python/mypy/pull/14041
 # see https://peps.python.org/pep-0673/#use-in-generic-classes
 # Self = t.TypeVar("Self", bound="TraitType[Any, Any]")
 if t.TYPE_CHECKING:
     from typing_extensions import Literal, Self
 
-    K = TypeVar("K", default=str)
-    V = TypeVar("V", default=t.Any)
 
 
 # We use a type for the getter (G) and setter (G) because we allow
@@ -519,7 +518,7 @@ class TraitType(BaseDescriptor, t.Generic[G, S]):
     default_value: t.Any = Undefined
 
     def __init__(
-        self: TraitType[G, S],
+        self,
         default_value: t.Any = Undefined,
         allow_none: bool = False,
         read_only: bool | None = None,
@@ -3580,7 +3579,7 @@ def item_from_string(self, s: str, index: int | None = None) -> T | str:
             return s
 
 
-class List(Container[t.List[T]]):
+class List(Container[t.List[T]], t.Generic[T]):
     """An instance of a Python list."""
 
     klass = list  # type:ignore[assignment]
@@ -3645,7 +3644,7 @@ def set(self, obj: t.Any, value: t.Any) -> None:
             return super().set(obj, value)
 
 
-class Set(Container[t.Set[t.Any]]):
+class Set(Container[t.Set[T]], t.Generic[T]):
     """An instance of a Python set."""
 
     klass = set
@@ -3656,8 +3655,8 @@ class Set(Container[t.Set[t.Any]]):
     # Redefine __init__ just to make the docstring more accurate.
     def __init__(
         self,
-        trait: t.Any = None,
-        default_value: t.Any = Undefined,
+        trait: TraitType[T, t.Any] | t.Iterable[T] | None = None,
+        default_value: set[T] | t.Any = Undefined,
         minlen: int = 0,
         maxlen: int = sys.maxsize,
         **kwargs: t.Any,
@@ -3720,13 +3719,13 @@ def default_value_repr(self) -> str:
         return "{" + list_repr[1:-1] + "}"
 
 
-class Tuple(Container[t.Tuple[t.Any, ...]]):
+class Tuple(Container[t.Tuple[T]], t.Generic[T]):
     """An instance of a Python tuple."""
 
     klass = tuple
     _cast_types = (list,)
 
-    def __init__(self, *traits: t.Any, **kwargs: t.Any) -> None:
+    def __init__(self, *traits: T, **kwargs: t.Any) -> None:
         """Create a tuple from a list, set, or tuple.
 
         Create a fixed-type tuple with Traits:
@@ -3849,7 +3848,7 @@ def subclass_init(self, cls: type[t.Any]) -> None:
         # to opt out of instance_init
 
 
-class Dict(Instance["dict[K, V]"]):
+class Dict(Instance["dict[K, V]"], t.Generic[K, V]):
     """An instance of a Python dict.
 
     One or more traits can be passed to the constructor
@@ -3869,9 +3868,9 @@ class Dict(Instance["dict[K, V]"]):
 
     def __init__(
         self,
-        value_trait: TraitType[t.Any, t.Any] | dict[K, V] | Sentinel | None = None,
+        value_trait: TraitType[V, t.Any] | dict[K, V] | Sentinel | None = None,
         per_key_traits: t.Any = None,
-        key_trait: TraitType[t.Any, t.Any] | None = None,
+        key_trait: TraitType[K, t.Any] | None = None,
         default_value: dict[K, V] | Sentinel | None = Undefined,
         **kwargs: t.Any,
     ) -> None:
@@ -4203,7 +4202,7 @@ def validate(self, obj: t.Any, value: t.Any) -> re.Pattern[t.Any] | None:
             self.error(obj, value)
 
 
-class UseEnum(TraitType[t.Any, t.Any]):
+class UseEnum(TraitType[G, S], t.Generic[G, S]):
     """Use a Enum class as model for the data type description.
     Note that if no default-value is provided, the first enum-value is used
     as default-value.
@@ -4236,9 +4235,7 @@ class MyEntity(HasTraits):
     default_value: enum.Enum | None = None
     info_text = "Trait type adapter to a Enum class"
 
-    def __init__(
-        self, enum_class: type[t.Any], default_value: t.Any = None, **kwargs: t.Any
-    ) -> None:
+    def __init__(self, enum_class: type[G], default_value: t.Any = None, **kwargs: t.Any) -> None:
         assert issubclass(enum_class, enum.Enum), "REQUIRE: enum.Enum, but was: %r" % enum_class
         allow_none = kwargs.get("allow_none", False)
         if default_value is None and not allow_none: