From 2c27a58ddf351d1433cafc9f62339b3b4899cc92 Mon Sep 17 00:00:00 2001 From: "Maarten A. Breddels" Date: Tue, 20 Dec 2022 14:08:00 +0100 Subject: [PATCH] implement TCPAddress, Int, CInt and Bool This makes use of Self, which only seems to work with the current mypy main branch. --- traitlets/tests/test_typing.py | 108 ++++++++++++- traitlets/traitlets.py | 281 +++++++++++++++++++++++++++------ 2 files changed, 337 insertions(+), 52 deletions(-) diff --git a/traitlets/tests/test_typing.py b/traitlets/tests/test_typing.py index bb54facd..4076067c 100644 --- a/traitlets/tests/test_typing.py +++ b/traitlets/tests/test_typing.py @@ -1,26 +1,120 @@ +from typing import Optional + import pytest from typing_extensions import reveal_type -from traitlets import Bool, HasTraits, TCPAddress +from traitlets import Bool, CInt, HasTraits, Int, TCPAddress @pytest.mark.mypy_testing def mypy_bool_typing(): class T(HasTraits): - b = Bool().tag(sync=True) + b = Bool(True).tag(sync=True) + ob = Bool(None, allow_none=True).tag(sync=True) t = T() - reveal_type(t.b) # R: Union[builtins.bool, None] + reveal_type( + Bool(True) + ) # R: traitlets.traitlets.Bool[builtins.bool, Union[builtins.bool, builtins.int]] + reveal_type( + Bool(True).tag(sync=True) + ) # R: traitlets.traitlets.Bool[builtins.bool, Union[builtins.bool, builtins.int]] + reveal_type( + Bool(None, allow_none=True) + ) # R: traitlets.traitlets.Bool[Union[builtins.bool, None], Union[builtins.bool, builtins.int, None]] + reveal_type( + Bool(None, allow_none=True).tag(sync=True) + ) # R: traitlets.traitlets.Bool[Union[builtins.bool, None], Union[builtins.bool, builtins.int, None]] + reveal_type( + T.b + ) # R: traitlets.traitlets.Bool[builtins.bool, Union[builtins.bool, builtins.int]] + reveal_type(t.b) # R: builtins.bool + reveal_type(t.ob) # R: Union[builtins.bool, None] + reveal_type( + T.b + ) # R: traitlets.traitlets.Bool[builtins.bool, Union[builtins.bool, builtins.int]] + reveal_type( + T.ob + ) # R: traitlets.traitlets.Bool[Union[builtins.bool, None], Union[builtins.bool, builtins.int, None]] # we would expect this to be Optional[Union[bool, int]], but... - t.b = "foo" # E: Incompatible types in assignment (expression has type "str", variable has type "Optional[int]") [assignment] - T.b.tag(foo=True) + t.b = "foo" # E: Incompatible types in assignment (expression has type "str", variable has type "Union[bool, int]") [assignment] + t.b = None # E: Incompatible types in assignment (expression has type "None", variable has type "Union[bool, int]") [assignment] + + +@pytest.mark.mypy_testing +def mypy_int_typing(): + class T(HasTraits): + i: Int[int, int] = Int(42).tag(sync=True) + oi: Int[Optional[int], Optional[int]] = Int(42, allow_none=True).tag(sync=True) + + t = T() + reveal_type(Int(True)) # R: traitlets.traitlets.Int[builtins.int, builtins.int] + reveal_type(Int(True).tag(sync=True)) # R: traitlets.traitlets.Int[builtins.int, builtins.int] + reveal_type( + Int(None, allow_none=True) + ) # R: traitlets.traitlets.Int[Union[builtins.int, None], Union[builtins.int, None]] + reveal_type( + Int(None, allow_none=True).tag(sync=True) + ) # R: traitlets.traitlets.Int[Union[builtins.int, None], Union[builtins.int, None]] + reveal_type(T.i) # R: traitlets.traitlets.Int[builtins.int, builtins.int] + reveal_type(t.i) # R: builtins.int + reveal_type(t.oi) # R: Union[builtins.int, None] + reveal_type(T.i) # R: traitlets.traitlets.Int[builtins.int, builtins.int] + reveal_type( + T.oi + ) # R: traitlets.traitlets.Int[Union[builtins.int, None], Union[builtins.int, None]] + t.i = "foo" # E: Incompatible types in assignment (expression has type "str", variable has type "int") [assignment] + t.i = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") [assignment] + t.i = 1.2 # E: Incompatible types in assignment (expression has type "float", variable has type "int") [assignment] + + +@pytest.mark.mypy_testing +def mypy_cint_typing(): + class T(HasTraits): + i = CInt(42).tag(sync=True) + oi = CInt(42, allow_none=True).tag(sync=True) + + t = T() + reveal_type(CInt(True)) # R: traitlets.traitlets.CInt[builtins.int, Union[builtins.int, Any]] + reveal_type( + CInt(True).tag(sync=True) + ) # R: traitlets.traitlets.CInt[builtins.int, Union[builtins.int, Any]] + reveal_type( + CInt(None, allow_none=True) + ) # R: traitlets.traitlets.CInt[Union[builtins.int, None], Union[builtins.int, Any, None]] + reveal_type( + CInt(None, allow_none=True).tag(sync=True) + ) # R: traitlets.traitlets.CInt[Union[builtins.int, None], Union[builtins.int, Any, None]] + reveal_type(T.i) # R: traitlets.traitlets.CInt[builtins.int, Union[builtins.int, Any]] + reveal_type(t.i) # R: builtins.int + reveal_type(t.oi) # R: Union[builtins.int, None] + reveal_type(T.i) # R: traitlets.traitlets.CInt[builtins.int, Union[builtins.int, Any]] + reveal_type( + T.oi + ) # R: traitlets.traitlets.CInt[Union[builtins.int, None], Union[builtins.int, Any, None]] @pytest.mark.mypy_testing def mypy_tcp_typing(): class T(HasTraits): tcp = TCPAddress() + otcp = TCPAddress(None, allow_none=True) t = T() - reveal_type(t.tcp) # R: Union[Tuple[builtins.str, builtins.int], None] - t.tcp = "foo" # E: Incompatible types in assignment (expression has type "str", variable has type "Optional[Tuple[str, int]]") [assignment] + reveal_type(t.tcp) # R: Tuple[builtins.str, builtins.int] + reveal_type( + T.tcp + ) # R: traitlets.traitlets.TCPAddress[Tuple[builtins.str, builtins.int], Tuple[builtins.str, builtins.int]] + reveal_type( + T.tcp.tag(sync=True) + ) # R:traitlets.traitlets.TCPAddress[Tuple[builtins.str, builtins.int], Tuple[builtins.str, builtins.int]] + reveal_type(t.otcp) # R: Union[Tuple[builtins.str, builtins.int], None] + reveal_type( + T.otcp + ) # R: traitlets.traitlets.TCPAddress[Union[Tuple[builtins.str, builtins.int], None], Union[Tuple[builtins.str, builtins.int], None]] + reveal_type( + T.otcp.tag(sync=True) + ) # R: traitlets.traitlets.TCPAddress[Union[Tuple[builtins.str, builtins.int], None], Union[Tuple[builtins.str, builtins.int], None]] + t.tcp = "foo" # E: Incompatible types in assignment (expression has type "str", variable has type "Tuple[str, int]") [assignment] + t.otcp = "foo" # E: Incompatible types in assignment (expression has type "str", variable has type "Optional[Tuple[str, int]]") [assignment] + t.tcp = None # E: Incompatible types in assignment (expression has type "None", variable has type "Tuple[str, int]") [assignment] diff --git a/traitlets/traitlets.py b/traitlets/traitlets.py index 58ca9c9a..b06be80e 100644 --- a/traitlets/traitlets.py +++ b/traitlets/traitlets.py @@ -526,7 +526,11 @@ def instance_init(self, obj): G = t.TypeVar("G") S = t.TypeVar("S") -# Self = t.TypeVar("Self", bound="TraitType") # Holdover waiting for typings.Self in Python 3.11 + +# 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]") +from typing_extensions import Self class TraitType(BaseDescriptor, t.Generic[G, S]): @@ -696,7 +700,56 @@ def get(self, obj: "HasTraits", cls: t.Any = None) -> t.Optional[G]: else: return value # type: ignore - def __get__(self, obj: "HasTraits", cls: t.Any = None) -> t.Optional[G]: + if t.TYPE_CHECKING: + # this gives ok type information, but not specific enought (e.g. it will) + # always be a TraitType, not a subclass, like Bool + + @t.overload + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[S, Sentinel] = Undefined, + allow_none: t.Literal[False] = False, + read_only: t.Optional[bool] = None, + help: t.Optional[str] = None, + config: t.Any = None, + **kwargs: t.Dict[str, t.Any], + ) -> "TraitType[G, S]": + ... + + @t.overload + def __new__( + cls, + default_value: t.Union[S, None, Sentinel] = Undefined, + allow_none: t.Literal[True] = True, + read_only: t.Optional[bool] = None, + help: t.Optional[str] = None, + config: t.Any = None, + **kwargs: t.Dict[str, t.Any], + ) -> "TraitType[t.Optional[G], S]": + ... + + def __new__( # type: ignore[no-untyped-def, misc] + cls, + default_value: t.Union[S, None, Sentinel] = Undefined, + allow_none: t.Literal[True, False] = False, + read_only=None, + help=None, + config=None, + **kwargs, + ) -> t.Union["TraitType[t.Optional[G], S]", "TraitType[G, S]"]: + ... + + if t.TYPE_CHECKING: + + @t.overload + def __get__(self, obj: None, cls: type[t.Any]) -> Self: + ... + + @t.overload + def __get__(self, obj: t.Any, cls: type[t.Any]) -> G: + ... + + def __get__(self, obj: t.Union["HasTraits", None], cls: type[t.Any]) -> t.Union[Self, G]: """Get the value of the trait by self.name for the instance. Default values are instantiated when :meth:`HasTraits.__new__` @@ -707,7 +760,7 @@ def __get__(self, obj: "HasTraits", cls: t.Any = None) -> t.Optional[G]: if obj is None: return self else: - return self.get(obj, cls) + return t.cast(G, self.get(obj, cls)) # the G should encode the Optional def set(self, obj, value): new_value = self._validate(obj, value) @@ -727,7 +780,7 @@ def set(self, obj, value): # comparison above returns something other than True/False obj._notify_trait(self.name, old_value, new_value) - def __set__(self, obj: "HasTraits", value: t.Optional[S]) -> None: + def __set__(self, obj: "HasTraits", value: S) -> None: """Set the value of the trait by self.name for the instance. Values pass through a validation stage where errors are raised when @@ -874,7 +927,7 @@ def set_metadata(self, key, value): warn("Deprecated in traitlets 4.1, " + msg, DeprecationWarning, stacklevel=2) self.metadata[key] = value - def tag(self, **metadata: t.Any) -> "TraitType[G, S]": + def tag(self, **metadata: t.Any) -> "Self": """Sets metadata and returns self. This allows convenient metadata tagging when initializing the trait, such as: @@ -1983,7 +2036,7 @@ def trait_events(cls, name=None): # ----------------------------------------------------------------------------- -class ClassBasedTraitType(TraitType[t.Any, t.Any]): +class ClassBasedTraitType(TraitType[G, S]): """ A trait with error reporting and string -> type resolution for Type, Instance and This. @@ -1996,7 +2049,7 @@ def _resolve_string(self, string): return import_item(string) -class Type(ClassBasedTraitType): +class Type(ClassBasedTraitType[G, S]): """A trait whose value must be a subclass of a specified class.""" def __init__(self, default_value=Undefined, klass=None, **kwargs): @@ -2095,7 +2148,7 @@ def default_value_repr(self): T = t.TypeVar("T") -class Instance(ClassBasedTraitType, t.Generic[T]): +class Instance(ClassBasedTraitType[T, T]): """A trait whose value must be an instance of a specified class. The value can also be an instance of a subclass of the specified class. @@ -2129,17 +2182,6 @@ def __new__( # type: ignore[misc] ) -> "Instance[T] | Instance[T | None]": ... - @t.overload # type:ignore[override] - def __get__(self, obj: None, cls: type[t.Any]) -> "Instance[T]": - ... - - @t.overload - def __get__(self, obj: t.Any, cls: type[t.Any]) -> T: - ... - - def __get__(self, obj: t.Union[t.Any, None], cls: type[t.Any]) -> "t.Union[T, Instance[T]]": - ... - klass: t.Optional[t.Union[str, t.Type[T]]] = None def __init__( @@ -2249,7 +2291,7 @@ def _resolve_string(self, string): return import_item(".".join([modname, string])) -class ForwardDeclaredType(ForwardDeclaredMixin, Type): +class ForwardDeclaredType(ForwardDeclaredMixin, Type[G, S]): """ Forward-declared version of Type. """ @@ -2265,7 +2307,7 @@ class ForwardDeclaredInstance(ForwardDeclaredMixin, Instance[T]): pass -class This(ClassBasedTraitType): +class This(ClassBasedTraitType[G, S]): """A trait for instances of the class containing this trait. Because how how and when class bodies are executed, the ``This`` @@ -2412,7 +2454,7 @@ def __new__( # type: ignore[misc] ) -> "Any": ... - @t.overload # type:ignore[override] + @t.overload def __get__(self, obj: None, cls: type[t.Any]) -> "Any": ... @@ -2457,16 +2499,60 @@ def _validate_bounds(trait, obj, value): return value -class Int(TraitType[int, int]): +class Int(TraitType[G, S]): """An int trait.""" default_value = 0 info_text = "an int" - def __init__(self, default_value=Undefined, allow_none=False, **kwargs): - self.min = kwargs.pop("min", None) - self.max = kwargs.pop("max", None) - super().__init__(default_value=default_value, allow_none=allow_none, **kwargs) + if t.TYPE_CHECKING: + + @t.overload + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[int, Sentinel] = ..., + allow_none: t.Literal[False] = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "Int[int, int]": + ... + + @t.overload + def __new__( + cls, + default_value: t.Union[int, None, Sentinel] = ..., + allow_none: t.Literal[True] = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "Int[t.Optional[int], t.Optional[int]]": + ... + + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[int, None, Sentinel] = Undefined, + allow_none: t.Literal[True, False] = False, + **kwargs: t.Dict[str, t.Any], + ) -> t.Union["Int[t.Optional[int], t.Optional[int]]", "Int[int, int]"]: + ... + + # if we use Self and a modern mypy, we don't need to do this + # if t.TYPE_CHECKING: + # @t.overload + # def __get__(self, obj: None, cls: type[t.Any]) -> "Int[G, S]": + # ... + + # @t.overload + # def __get__(self, obj: t.Any, cls: type[t.Any]) -> G: + # ... + + # def __get__(self, obj: t.Union[t.Any, None], cls: type[t.Any]) -> t.Union[G, "Int[G, S]"]: + # ... + + # mypy infers Int[, ] if this constructor exists + if not t.TYPE_CHECKING: + + def __init__(self, default_value=Undefined, allow_none=False, **kwargs): + self.min = kwargs.pop("min", None) + self.max = kwargs.pop("max", None) + super().__init__(default_value=default_value, allow_none=allow_none, **kwargs) def validate(self, obj, value): if not isinstance(value, int): @@ -2482,9 +2568,48 @@ def subclass_init(self, cls): pass # fully opt out of instance_init -class CInt(Int, TraitType[int, t.Any]): +class CInt(Int[G, S]): """A casting version of the int trait.""" + if t.TYPE_CHECKING: + + @t.overload + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[int, Sentinel] = ..., + allow_none: t.Literal[False] = False, + read_only: t.Optional[bool] = ..., + help: t.Optional[str] = ..., + config: t.Any = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "CInt[int, t.Union[int, t.Any]]": + ... + + @t.overload + def __new__( + cls, + default_value: t.Union[int, None, Sentinel] = ..., + allow_none: t.Literal[True] = True, + read_only: t.Optional[bool] = ..., + help: t.Optional[str] = ..., + config: t.Any = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "CInt[t.Optional[int], t.Union[int, t.Any, None]]": + ... + + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[bool, None, Sentinel] = Undefined, + allow_none: t.Literal[True, False] = False, + read_only: t.Optional[bool] = None, + help: t.Optional[str] = None, + config: t.Any = None, + **kwargs: t.Dict[str, t.Any], + ) -> t.Union[ + "CInt[t.Optional[int], t.Union[int, t.Any, None]]", "CInt[int, t.Union[int, t.Any]]" + ]: + ... + def validate(self, obj, value): try: value = int(value) @@ -2649,7 +2774,7 @@ def __new__( # type: ignore[misc] ) -> "Unicode": ... - @t.overload # type:ignore[override] + @t.overload def __get__(self, obj: None, cls: type[t.Any]) -> "Unicode": ... @@ -2735,12 +2860,51 @@ def validate(self, obj, value): self.error(obj, value) -class Bool(TraitType[bool, t.Union[bool, int]]): +class Bool(TraitType[G, S]): """A boolean (True, False) trait.""" default_value = False info_text = "a boolean" + if t.TYPE_CHECKING: + + @t.overload + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[bool, Sentinel] = ..., + allow_none: t.Literal[False] = False, + read_only: t.Optional[bool] = ..., + help: t.Optional[str] = ..., + config: t.Any = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "Bool[bool, t.Union[bool, int]]": + ... + + @t.overload + def __new__( + cls, + default_value: t.Union[bool, None, Sentinel] = ..., + allow_none: t.Literal[True] = True, + read_only: t.Optional[bool] = ..., + help: t.Optional[str] = ..., + config: t.Any = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "Bool[t.Optional[bool], t.Union[bool, int, None]]": + ... + + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[bool, None, Sentinel] = Undefined, + allow_none: t.Literal[True, False] = False, + read_only: t.Optional[bool] = None, + help: t.Optional[str] = None, + config: t.Any = None, + **kwargs: t.Dict[str, t.Any], + ) -> t.Union[ + "Bool[t.Optional[bool], t.Union[bool, int, None]]", "Bool[bool, t.Union[bool, int]]" + ]: + ... + def validate(self, obj, value): if isinstance(value, bool): return value @@ -2773,7 +2937,7 @@ def argcompleter(self, **kwargs): return completions -class CBool(Bool, TraitType[bool, t.Any]): +class CBool(Bool[G, S]): """A casting version of the boolean trait.""" def validate(self, obj, value): @@ -2944,19 +3108,6 @@ def __new__( # type: ignore[misc] ) -> "Container[T] | Container[T | None]": ... - @t.overload # type:ignore[override] - def __get__(self, obj: None, cls: type[t.Any]) -> "Container[T]": - ... - - @t.overload - def __get__(self, obj: t.Any, cls: type[t.Any]) -> T: - ... - - def __get__( - self, obj: t.Union[t.Any, None], cls: type[t.Any] - ) -> "t.Union[T, Container[T]]": - ... - klass: t.Optional[t.Type[T]] = None _cast_types: t.Any = () _valid_defaults = SequenceTypes @@ -3657,7 +3808,7 @@ def item_from_string(self, s): return {key: value} -class TCPAddress(TraitType[tuple[str, int], tuple[str, int]]): +class TCPAddress(TraitType[G, S]): """A trait for an (ip, port) tuple. This allows for both IPv4 IP addresses as well as hostnames. @@ -3666,6 +3817,46 @@ class TCPAddress(TraitType[tuple[str, int], tuple[str, int]]): default_value = ("127.0.0.1", 0) info_text = "an (ip, port) tuple" + if t.TYPE_CHECKING: + + @t.overload + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[bool, Sentinel] = ..., + allow_none: t.Literal[False] = False, + read_only: t.Optional[bool] = ..., + help: t.Optional[str] = ..., + config: t.Any = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "TCPAddress[tuple[str, int], tuple[str, int]]": + ... + + @t.overload + def __new__( + cls, + default_value: t.Union[bool, None, Sentinel] = ..., + allow_none: t.Literal[True] = True, + read_only: t.Optional[bool] = ..., + help: t.Optional[str] = ..., + config: t.Any = ..., + **kwargs: t.Dict[str, t.Any], + ) -> "TCPAddress[t.Optional[tuple[str, int]], t.Optional[tuple[str, int]]]": + ... + + def __new__( # type: ignore[misc] + cls, + default_value: t.Union[bool, None, Sentinel] = Undefined, + allow_none: t.Literal[True, False] = False, + read_only: t.Optional[bool] = None, + help: t.Optional[str] = None, + config: t.Any = None, + **kwargs: t.Dict[str, t.Any], + ) -> t.Union[ + "TCPAddress[t.Optional[tuple[str, int]], t.Optional[tuple[str, int]]]", + "TCPAddress[tuple[str, int], tuple[str, int]]", + ]: + ... + def validate(self, obj, value): if isinstance(value, tuple): if len(value) == 2: