Improve Int and Float type stubs (#1656)

This PR improves the type stubs for Int and Float, being more accurate about what's accepted for those trait types.

We need SupportsIndex, which is only available in typing from Python 3.8 onwards. So I've added typing-extensions as a dependency for traits-stubs. We can drop the dependency once we no longer support Python 3.6 and Python 3.7.

traits-stubs/setup.py

@@ -52,7 +52,12 @@ def get_long_description():
         long_description=get_long_description(),
         long_description_content_type="text/x-rst",
         download_url="https://pypi.python.org/pypi/traits-stubs",
-        install_requires=["traits"],
+        install_requires=[
+            "traits",
+            # We need typing-extensions for SupportsIndex; once we no longer
+            # support Python < 3.8, we can drop this requirement.
+            'typing-extensions; python_version<"3.8"',
+        ],
         extras_require={"test": ["mypy"]},
         packages=[
             "traits-stubs",

traits-stubs/traits-stubs/trait_types.pyi

@@ -10,6 +10,7 @@
 
 import datetime
 from pathlib import PurePath as _PurePath
+import sys
 from typing import (
     Any as _Any,
     Callable as _CallableType,
@@ -18,12 +19,20 @@ from typing import (
     Optional,
     Sequence as _Sequence,
     Set as _SetType,
+    SupportsFloat,
     Type as _Type,
     TypeVar,
     Union as _Union,
 )
 from uuid import UUID as _UUID
 
+# Once we no longer support Python 3.6 or Python 3.7, we can import
+# SupportsIndex from typing instead of typing_extensions.
+if sys.version_info < (3, 8):
+    from typing_extensions import SupportsIndex
+else:
+    from typing import SupportsIndex
+
 from .trait_type import _TraitType
 
 SetTypes: _Any
@@ -60,7 +69,7 @@ class _BaseInt(_TraitType[_T, int]):
     ...
 
 
-class BaseInt(_BaseInt[int]):
+class BaseInt(_TraitType[SupportsIndex, int]):
     ...
 
 
@@ -72,7 +81,7 @@ class _BaseFloat(_TraitType[_T, float]):
     ...
 
 
-class BaseFloat(_BaseFloat[float]):
+class BaseFloat(_TraitType[_Union[SupportsFloat, SupportsIndex], float]):
     ...
 
 

traits-stubs/traits_stubs_tests/examples/Float.py

@@ -11,6 +11,18 @@
 from traits.api import Float, HasTraits
 
 
+class HasIndex:
+    """Class with __index__ method; instances should be assignable to Float."""
+    def __index__(self):
+        return 1729
+
+
+class HasFloat:
+    """Class with __float__ method; instances should be assignable to Float."""
+    def __float__(self):
+        return 1729.0
+
+
 class Test(HasTraits):
     i = Float()
 
@@ -19,3 +31,5 @@ class Test(HasTraits):
 o.i = "5"  # E: assignment
 o.i = 5
 o.i = 5.5
+o.i = HasIndex()
+o.i = HasFloat()

traits-stubs/traits_stubs_tests/examples/Int.py

@@ -11,6 +11,12 @@
 from traits.api import HasTraits, Int
 
 
+class HasIndex:
+    """Class with __index__ method; instances should be assignable to Int."""
+    def __index__(self):
+        return 1729
+
+
 class Test(HasTraits):
     i = Int()
     j = Int(default_value="234")  # E: arg-type
@@ -19,6 +25,7 @@ class Test(HasTraits):
 
 o = Test()
 o.i = 5
+o.i = HasIndex()
 
 o.i = "5"  # E: assignment
 o.i = 5.5  # E: assignment