Merge pull request #18322 from BvB93/plugin-precision

ENH: Add a mypy plugin for exposing platform-specific extended-precision types

8 files changed, 256 insertions, 37 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 5f731683d..19711d80b 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -129,6 +129,23 @@ from numpy.typing._callable import (
     _ComparisonOp,
 )
 
+# NOTE: Numpy's mypy plugin is used for removing the types unavailable
+# to the specific platform
+from numpy.typing._extended_precision import (
+    uint128 as uint128,
+    uint256 as uint256,
+    int128 as int128,
+    int256 as int256,
+    float80 as float80,
+    float96 as float96,
+    float128 as float128,
+    float256 as float256,
+    complex160 as complex160,
+    complex192 as complex192,
+    complex256 as complex256,
+    complex512 as complex512,
+)
+
 from typing import (
     Any,
     ByteString,
@@ -2507,7 +2524,6 @@ class floating(inexact[_NBit1]):
 float16 = floating[_16Bit]
 float32 = floating[_32Bit]
 float64 = floating[_64Bit]
-float128 = floating[_128Bit]
 
 half = floating[_NBitHalf]
 single = floating[_NBitSingle]
@@ -2542,7 +2558,6 @@ class complexfloating(inexact[_NBit1], Generic[_NBit1, _NBit2]):
 
 complex64 = complexfloating[_32Bit, _32Bit]
 complex128 = complexfloating[_64Bit, _64Bit]
-complex256 = complexfloating[_128Bit, _128Bit]
 
 csingle = complexfloating[_NBitSingle, _NBitSingle]
 singlecomplex = complexfloating[_NBitSingle, _NBitSingle]
@@ -2596,8 +2611,6 @@ class str_(character, str):
 unicode_ = str_
 str0 = str_
 
-# TODO: Platform dependent types: float128, complex256, float96
-
 def array(
     object: object,
     dtype: DTypeLike = ...,
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
index 8147789fb..8f5df483b 100644
--- a/numpy/typing/__init__.py
+++ b/numpy/typing/__init__.py
@@ -22,14 +22,19 @@ the two below:
 Mypy plugin
 -----------
 
-A mypy_ plugin is available for automatically assigning the (platform-dependent)
-precisions of certain `~numpy.number` subclasses, including the likes of
-`~numpy.int_`, `~numpy.intp` and `~numpy.longlong`. See the documentation on
-:ref:`scalar types <arrays.scalars.built-in>` for a comprehensive overview
-of the affected classes.
-
-Note that while usage of the plugin is completely optional, without it the
-precision of above-mentioned classes will be inferred as `~typing.Any`.
+A mypy_ plugin is distributed in `numpy.typing` for managing a number of
+platform-specific annotations. Its function can be split into to parts:
+
+* Assigning the (platform-dependent) precisions of certain `~numpy.number` subclasses,
+  including the likes of `~numpy.int_`, `~numpy.intp` and `~numpy.longlong`.
+  See the documentation on :ref:`scalar types <arrays.scalars.built-in>` for a
+  comprehensive overview of the affected classes. without the plugin the precision
+  of all relevant classes will be inferred as `~typing.Any`.
+* Removing all extended-precision `~numpy.number` subclasses that are unavailable
+  for the platform in question. Most notable this includes the likes of
+  `~numpy.float128` and `~numpy.complex256`. Without the plugin *all*
+  extended-precision types will, as far as mypy is concerned, be available
+  to all platforms.
 
 To enable the plugin, one must add it to their mypy `configuration file`_:
 
diff --git a/numpy/typing/_callable.py b/numpy/typing/_callable.py
index 0e34cd799..23cabd47a 100644
--- a/numpy/typing/_callable.py
+++ b/numpy/typing/_callable.py
@@ -38,7 +38,7 @@ from numpy import (
     complexfloating,
     complex128,
 )
-from ._nbit import _NBitInt
+from ._nbit import _NBitInt, _NBitDouble
 from ._scalars import (
     _BoolLike_co,
     _IntLike_co,
@@ -133,7 +133,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload  # platform dependent
         def __call__(self, __other: int) -> _2Tuple[int_]: ...
         @overload
-        def __call__(self, __other: float) -> _2Tuple[float64]: ...
+        def __call__(self, __other: float) -> _2Tuple[floating[Union[_NBit1, _NBitDouble]]]: ...
         @overload
         def __call__(self, __other: _IntType) -> _2Tuple[_IntType]: ...
         @overload
@@ -151,11 +151,13 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: bool) -> floating[_NBit1]: ...
         @overload
-        def __call__(self, __other: int) -> floating[_NBitInt]: ...
+        def __call__(self, __other: int) -> floating[Union[_NBit1, _NBitInt]]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
-        def __call__(self, __other: complex) -> complex128: ...
+        def __call__(
+            self, __other: complex
+        ) -> complexfloating[Union[_NBit1, _NBitDouble], Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(self, __other: integer[_NBit2]) -> floating[Union[_NBit1, _NBit2]]: ...
 
@@ -168,9 +170,11 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
             self, __other: Union[int, signedinteger[Any]]
         ) -> Union[signedinteger[Any], float64]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
-        def __call__(self, __other: complex) -> complex128: ...
+        def __call__(
+            self, __other: complex
+        ) -> complexfloating[Union[_NBit1, _NBitDouble], Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self, __other: unsignedinteger[_NBit2]
@@ -196,7 +200,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
             self, __other: Union[int, signedinteger[Any]]
         ) -> Union[signedinteger[Any], float64]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self, __other: unsignedinteger[_NBit2]
@@ -210,7 +214,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
             self, __other: Union[int, signedinteger[Any]]
         ) -> Union[_2Tuple[signedinteger[Any]], _2Tuple[float64]]: ...
         @overload
-        def __call__(self, __other: float) -> _2Tuple[float64]: ...
+        def __call__(self, __other: float) -> _2Tuple[floating[Union[_NBit1, _NBitDouble]]]: ...
         @overload
         def __call__(
             self, __other: unsignedinteger[_NBit2]
@@ -222,9 +226,11 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> signedinteger[Union[_NBit1, _NBitInt]]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
-        def __call__(self, __other: complex) -> complex128: ...
+        def __call__(
+            self, __other: complex
+        ) -> complexfloating[Union[_NBit1, _NBitDouble], Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self, __other: signedinteger[_NBit2]
@@ -246,7 +252,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> signedinteger[Union[_NBit1, _NBitInt]]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self, __other: signedinteger[_NBit2]
@@ -258,7 +264,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> _2Tuple[signedinteger[Union[_NBit1, _NBitInt]]]: ...
         @overload
-        def __call__(self, __other: float) -> _2Tuple[float64]: ...
+        def __call__(self, __other: float) -> _2Tuple[floating[Union[_NBit1, _NBitDouble]]]: ...
         @overload
         def __call__(
             self, __other: signedinteger[_NBit2]
@@ -270,9 +276,11 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> floating[Union[_NBit1, _NBitInt]]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
-        def __call__(self, __other: complex) -> complex128: ...
+        def __call__(
+            self, __other: complex
+        ) -> complexfloating[Union[_NBit1, _NBitDouble], Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self, __other: Union[integer[_NBit2], floating[_NBit2]]
@@ -284,7 +292,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> floating[Union[_NBit1, _NBitInt]]: ...
         @overload
-        def __call__(self, __other: float) -> float64: ...
+        def __call__(self, __other: float) -> floating[Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self, __other: Union[integer[_NBit2], floating[_NBit2]]
@@ -296,7 +304,7 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> _2Tuple[floating[Union[_NBit1, _NBitInt]]]: ...
         @overload
-        def __call__(self, __other: float) -> _2Tuple[float64]: ...
+        def __call__(self, __other: float) -> _2Tuple[floating[Union[_NBit1, _NBitDouble]]]: ...
         @overload
         def __call__(
             self, __other: Union[integer[_NBit2], floating[_NBit2]]
@@ -308,7 +316,9 @@ if TYPE_CHECKING or HAVE_PROTOCOL:
         @overload
         def __call__(self, __other: int) -> complexfloating[Union[_NBit1, _NBitInt], Union[_NBit1, _NBitInt]]: ...
         @overload
-        def __call__(self, __other: Union[float, complex]) -> complex128: ...
+        def __call__(
+            self, __other: Union[float, complex]
+        ) -> complexfloating[Union[_NBit1, _NBitDouble], Union[_NBit1, _NBitDouble]]: ...
         @overload
         def __call__(
             self,
diff --git a/numpy/typing/_extended_precision.py b/numpy/typing/_extended_precision.py
new file mode 100644
index 000000000..3f1ce2038
--- /dev/null
+++ b/numpy/typing/_extended_precision.py
@@ -0,0 +1,42 @@
+"""A module with platform-specific extended precision `numpy.number` subclasses.
+
+The subclasses are defined here (instead of ``__init__.pyi``) such
+that they can be imported conditionally via the numpy's mypy plugin.
+"""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from . import (
+    _80Bit,
+    _96Bit,
+    _128Bit,
+    _256Bit,
+)
+
+if TYPE_CHECKING:
+    uint128 = np.unsignedinteger[_128Bit]
+    uint256 = np.unsignedinteger[_256Bit]
+    int128 = np.signedinteger[_128Bit]
+    int256 = np.signedinteger[_256Bit]
+    float80 = np.floating[_80Bit]
+    float96 = np.floating[_96Bit]
+    float128 = np.floating[_128Bit]
+    float256 = np.floating[_256Bit]
+    complex160 = np.complexfloating[_80Bit, _80Bit]
+    complex192 = np.complexfloating[_96Bit, _96Bit]
+    complex256 = np.complexfloating[_128Bit, _128Bit]
+    complex512 = np.complexfloating[_256Bit, _256Bit]
+else:
+    uint128 = NotImplemented
+    uint256 = NotImplemented
+    int128 = NotImplemented
+    int256 = NotImplemented
+    float80 = NotImplemented
+    float96 = NotImplemented
+    float128 = NotImplemented
+    float256 = NotImplemented
+    complex160 = NotImplemented
+    complex192 = NotImplemented
+    complex256 = NotImplemented
+    complex512 = NotImplemented
diff --git a/numpy/typing/mypy_plugin.py b/numpy/typing/mypy_plugin.py
index bdd5c50f3..901bf4fb1 100644
--- a/numpy/typing/mypy_plugin.py
+++ b/numpy/typing/mypy_plugin.py
@@ -1,5 +1,7 @@
 """A module containing `numpy`-specific plugins for mypy."""
 
+from __future__ import annotations
+
 import typing as t
 
 import numpy as np
@@ -8,6 +10,9 @@ try:
     import mypy.types
     from mypy.types import Type
     from mypy.plugin import Plugin, AnalyzeTypeContext
+    from mypy.nodes import MypyFile, ImportFrom, Statement
+    from mypy.build import PRI_MED
+
     _HookFunc = t.Callable[[AnalyzeTypeContext], Type]
     MYPY_EX: t.Optional[ModuleNotFoundError] = None
 except ModuleNotFoundError as ex:
@@ -37,12 +42,34 @@ def _get_precision_dict() -> t.Dict[str, str]:
     return ret
 
 
+def _get_extended_precision_list() -> t.List[str]:
+    extended_types = [np.ulonglong, np.longlong, np.longdouble, np.clongdouble]
+    extended_names = {
+        "uint128",
+        "uint256",
+        "int128",
+        "int256",
+        "float80",
+        "float96",
+        "float128",
+        "float256",
+        "complex160",
+        "complex192",
+        "complex256",
+        "complex512",
+    }
+    return [i.__name__ for i in extended_types if i.__name__ in extended_names]
+
+
 #: A dictionary mapping type-aliases in `numpy.typing._nbit` to
 #: concrete `numpy.typing.NBitBase` subclasses.
-_PRECISION_DICT = _get_precision_dict()
+_PRECISION_DICT: t.Final = _get_precision_dict()
+
+#: A list with the names of all extended precision `np.number` subclasses.
+_EXTENDED_PRECISION_LIST: t.Final = _get_extended_precision_list()
 
 
-def _hook(ctx: "AnalyzeTypeContext") -> "Type":
+def _hook(ctx: AnalyzeTypeContext) -> Type:
     """Replace a type-alias with a concrete ``NBitBase`` subclass."""
     typ, _, api = ctx
     name = typ.name.split(".")[-1]
@@ -50,20 +77,55 @@ def _hook(ctx: "AnalyzeTypeContext") -> "Type":
     return api.named_type(name_new)
 
 
-if MYPY_EX is None:
+if t.TYPE_CHECKING or MYPY_EX is None:
+    def _index(iterable: t.Iterable[Statement], id: str) -> int:
+        """Identify the first ``ImportFrom`` instance the specified `id`."""
+        for i, value in enumerate(iterable):
+            if getattr(value, "id", None) == id:
+                return i
+        else:
+            raise ValueError("Failed to identify a `ImportFrom` instance "
+                             f"with the following id: {id!r}")
+
     class _NumpyPlugin(Plugin):
         """A plugin for assigning platform-specific `numpy.number` precisions."""
 
         def get_type_analyze_hook(self, fullname: str) -> t.Optional[_HookFunc]:
+            """Set the precision of platform-specific `numpy.number` subclasses.
+
+            For example: `numpy.int_`, `numpy.longlong` and `numpy.longdouble`.
+            """
             if fullname in _PRECISION_DICT:
                 return _hook
             return None
 
+        def get_additional_deps(self, file: MypyFile) -> t.List[t.Tuple[int, str, int]]:
+            """Import platform-specific extended-precision `numpy.number` subclasses.
+
+            For example: `numpy.float96`, `numpy.float128` and `numpy.complex256`.
+            """
+            ret = [(PRI_MED, file.fullname, -1)]
+            if file.fullname == "numpy":
+                # Import ONLY the extended precision types available to the
+                # platform in question
+                imports = ImportFrom(
+                    "numpy.typing._extended_precision", 0,
+                    names=[(v, v) for v in _EXTENDED_PRECISION_LIST],
+                )
+                imports.is_top_level = True
+
+                # Replace the much broader extended-precision import
+                # (defined in `numpy/__init__.pyi`) with a more specific one
+                for lst in [file.defs, file.imports]:  # type: t.List[Statement]
+                    i = _index(lst, "numpy.typing._extended_precision")
+                    lst[i] = imports
+            return ret
+
     def plugin(version: str) -> t.Type[_NumpyPlugin]:
         """An entry-point for mypy."""
         return _NumpyPlugin
 
 else:
-    def plugin(version: str) -> t.Type["_NumpyPlugin"]:
+    def plugin(version: str) -> t.Type[_NumpyPlugin]:
         """An entry-point for mypy."""
         raise MYPY_EX
diff --git a/numpy/typing/tests/data/misc/extended_precision.py b/numpy/typing/tests/data/misc/extended_precision.py
new file mode 100644
index 000000000..1e495e4f3
--- /dev/null
+++ b/numpy/typing/tests/data/misc/extended_precision.py
@@ -0,0 +1,17 @@
+import numpy as np
+
+reveal_type(np.uint128())
+reveal_type(np.uint256())
+
+reveal_type(np.int128())
+reveal_type(np.int256())
+
+reveal_type(np.float80())
+reveal_type(np.float96())
+reveal_type(np.float128())
+reveal_type(np.float256())
+
+reveal_type(np.complex160())
+reveal_type(np.complex192())
+reveal_type(np.complex256())
+reveal_type(np.complex512())
diff --git a/numpy/typing/tests/data/reveal/arithmetic.py b/numpy/typing/tests/data/reveal/arithmetic.py
index b45330910..53a3a846d 100644
--- a/numpy/typing/tests/data/reveal/arithmetic.py
+++ b/numpy/typing/tests/data/reveal/arithmetic.py
@@ -1,5 +1,9 @@
 from typing import Any, List
 import numpy as np
+import numpy.typing as npt
+
+# Can't directly import `np.float128` as it is not available on all platforms
+f16: np.floating[npt._128Bit]
 
 c16 = np.complex128()
 f8 = np.float64()
@@ -263,6 +267,7 @@ reveal_type(AR_LIKE_O // AR_O)  # E: Any
 
 # unary ops
 
+reveal_type(-f16)  # E: {float128}
 reveal_type(-c16)  # E: {complex128}
 reveal_type(-c8)  # E: {complex64}
 reveal_type(-f8)  # E: {float64}
@@ -274,6 +279,7 @@ reveal_type(-u4)  # E: {uint32}
 reveal_type(-td)  # E: numpy.timedelta64
 reveal_type(-AR_f)  # E: Any
 
+reveal_type(+f16)  # E: {float128}
 reveal_type(+c16)  # E: {complex128}
 reveal_type(+c8)  # E: {complex64}
 reveal_type(+f8)  # E: {float64}
@@ -285,6 +291,7 @@ reveal_type(+u4)  # E: {uint32}
 reveal_type(+td)  # E: numpy.timedelta64
 reveal_type(+AR_f)  # E: Any
 
+reveal_type(abs(f16))  # E: {float128}
 reveal_type(abs(c16))  # E: {float64}
 reveal_type(abs(c8))  # E: {float32}
 reveal_type(abs(f8))  # E: {float64}
@@ -332,6 +339,7 @@ reveal_type(b_ / i4)  # E: {float64}
 reveal_type(b_ / u8)  # E: {float64}
 reveal_type(b_ / u4)  # E: {float64}
 reveal_type(b_ / f)  # E: {float64}
+reveal_type(b_ / f16)  # E: {float128}
 reveal_type(b_ / f8)  # E: {float64}
 reveal_type(b_ / f4)  # E: {float32}
 reveal_type(b_ / c)  # E: {complex128}
@@ -346,6 +354,7 @@ reveal_type(i4 / b_)  # E: {float64}
 reveal_type(u8 / b_)  # E: {float64}
 reveal_type(u4 / b_)  # E: {float64}
 reveal_type(f / b_)  # E: {float64}
+reveal_type(f16 / b_)  # E: {float128}
 reveal_type(f8 / b_)  # E: {float64}
 reveal_type(f4 / b_)  # E: {float32}
 reveal_type(c / b_)  # E: {complex128}
@@ -354,6 +363,7 @@ reveal_type(c8 / b_)  # E: {complex64}
 
 # Complex
 
+reveal_type(c16 + f16)  # E: {complex256}
 reveal_type(c16 + c16)  # E: {complex128}
 reveal_type(c16 + f8)  # E: {complex128}
 reveal_type(c16 + i8)  # E: {complex128}
@@ -364,10 +374,10 @@ reveal_type(c16 + b_)  # E: {complex128}
 reveal_type(c16 + b)  # E: {complex128}
 reveal_type(c16 + c)  # E: {complex128}
 reveal_type(c16 + f)  # E: {complex128}
-
 reveal_type(c16 + i)  # E: {complex128}
 reveal_type(c16 + AR_f)  # E: Any
 
+reveal_type(f16 + c16)  # E: {complex256}
 reveal_type(c16 + c16)  # E: {complex128}
 reveal_type(f8 + c16)  # E: {complex128}
 reveal_type(i8 + c16)  # E: {complex128}
@@ -381,6 +391,7 @@ reveal_type(f + c16)  # E: {complex128}
 reveal_type(i + c16)  # E: {complex128}
 reveal_type(AR_f + c16)  # E: Any
 
+reveal_type(c8 + f16)  # E: {complex256}
 reveal_type(c8 + c16)  # E: {complex128}
 reveal_type(c8 + f8)  # E: {complex128}
 reveal_type(c8 + i8)  # E: {complex128}
@@ -394,6 +405,7 @@ reveal_type(c8 + f)  # E: {complex128}
 reveal_type(c8 + i)  # E: numpy.complexfloating[{_NBitInt}, {_NBitInt}]
 reveal_type(c8 + AR_f)  # E: Any
 
+reveal_type(f16 + c8)  # E: {complex256}
 reveal_type(c16 + c8)  # E: {complex128}
 reveal_type(f8 + c8)  # E: {complex128}
 reveal_type(i8 + c8)  # E: {complex128}
@@ -409,6 +421,7 @@ reveal_type(AR_f + c8)  # E: Any
 
 # Float
 
+reveal_type(f8 + f16)  # E: {float128}
 reveal_type(f8 + f8)  # E: {float64}
 reveal_type(f8 + i8)  # E: {float64}
 reveal_type(f8 + f4)  # E: {float64}
@@ -420,6 +433,7 @@ reveal_type(f8 + f)  # E: {float64}
 reveal_type(f8 + i)  # E: {float64}
 reveal_type(f8 + AR_f)  # E: Any
 
+reveal_type(f16 + f8)  # E: {float128}
 reveal_type(f8 + f8)  # E: {float64}
 reveal_type(i8 + f8)  # E: {float64}
 reveal_type(f4 + f8)  # E: {float64}
@@ -431,6 +445,7 @@ reveal_type(f + f8)  # E: {float64}
 reveal_type(i + f8)  # E: {float64}
 reveal_type(AR_f + f8)  # E: Any
 
+reveal_type(f4 + f16)  # E: {float128}
 reveal_type(f4 + f8)  # E: {float64}
 reveal_type(f4 + i8)  # E: {float64}
 reveal_type(f4 + f4)  # E: {float32}
@@ -442,6 +457,7 @@ reveal_type(f4 + f)  # E: {float64}
 reveal_type(f4 + i)  # E: numpy.floating[{_NBitInt}]
 reveal_type(f4 + AR_f)  # E: Any
 
+reveal_type(f16 + f4)  # E: {float128}
 reveal_type(f8 + f4)  # E: {float64}
 reveal_type(i8 + f4)  # E: {float64}
 reveal_type(f4 + f4)  # E: {float32}
diff --git a/numpy/typing/tests/test_typing.py b/numpy/typing/tests/test_typing.py
index fc42e87b1..eb7e0b09e 100644
--- a/numpy/typing/tests/test_typing.py
+++ b/numpy/typing/tests/test_typing.py
@@ -8,7 +8,7 @@ from typing import Optional, IO, Dict, List
 
 import pytest
 import numpy as np
-from numpy.typing.mypy_plugin import _PRECISION_DICT
+from numpy.typing.mypy_plugin import _PRECISION_DICT, _EXTENDED_PRECISION_LIST
 
 try:
     from mypy import api
@@ -22,6 +22,7 @@ DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
 PASS_DIR = os.path.join(DATA_DIR, "pass")
 FAIL_DIR = os.path.join(DATA_DIR, "fail")
 REVEAL_DIR = os.path.join(DATA_DIR, "reveal")
+MISC_DIR = os.path.join(DATA_DIR, "misc")
 MYPY_INI = os.path.join(DATA_DIR, "mypy.ini")
 CACHE_DIR = os.path.join(DATA_DIR, ".mypy_cache")
 
@@ -51,7 +52,7 @@ def run_mypy() -> None:
     if os.path.isdir(CACHE_DIR):
         shutil.rmtree(CACHE_DIR)
 
-    for directory in (PASS_DIR, REVEAL_DIR, FAIL_DIR):
+    for directory in (PASS_DIR, REVEAL_DIR, FAIL_DIR, MISC_DIR):
         # Run mypy
         stdout, stderr, _ = api.run([
             "--config-file",
@@ -158,15 +159,27 @@ def _construct_format_dict():
         "uint16": "numpy.unsignedinteger[numpy.typing._16Bit]",
         "uint32": "numpy.unsignedinteger[numpy.typing._32Bit]",
         "uint64": "numpy.unsignedinteger[numpy.typing._64Bit]",
+        "uint128": "numpy.unsignedinteger[numpy.typing._128Bit]",
+        "uint256": "numpy.unsignedinteger[numpy.typing._256Bit]",
         "int8": "numpy.signedinteger[numpy.typing._8Bit]",
         "int16": "numpy.signedinteger[numpy.typing._16Bit]",
         "int32": "numpy.signedinteger[numpy.typing._32Bit]",
         "int64": "numpy.signedinteger[numpy.typing._64Bit]",
+        "int128": "numpy.signedinteger[numpy.typing._128Bit]",
+        "int256": "numpy.signedinteger[numpy.typing._256Bit]",
         "float16": "numpy.floating[numpy.typing._16Bit]",
         "float32": "numpy.floating[numpy.typing._32Bit]",
         "float64": "numpy.floating[numpy.typing._64Bit]",
+        "float80": "numpy.floating[numpy.typing._80Bit]",
+        "float96": "numpy.floating[numpy.typing._96Bit]",
+        "float128": "numpy.floating[numpy.typing._128Bit]",
+        "float256": "numpy.floating[numpy.typing._256Bit]",
         "complex64": "numpy.complexfloating[numpy.typing._32Bit, numpy.typing._32Bit]",
         "complex128": "numpy.complexfloating[numpy.typing._64Bit, numpy.typing._64Bit]",
+        "complex160": "numpy.complexfloating[numpy.typing._80Bit, numpy.typing._80Bit]",
+        "complex192": "numpy.complexfloating[numpy.typing._96Bit, numpy.typing._96Bit]",
+        "complex256": "numpy.complexfloating[numpy.typing._128Bit, numpy.typing._128Bit]",
+        "complex512": "numpy.complexfloating[numpy.typing._256Bit, numpy.typing._256Bit]",
 
         "ubyte": f"numpy.unsignedinteger[{dct['_NBitByte']}]",
         "ushort": f"numpy.unsignedinteger[{dct['_NBitShort']}]",
@@ -268,3 +281,44 @@ def test_code_runs(path):
     spec = importlib.util.spec_from_file_location(f"{dirname}.{filename}", path)
     test_module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(test_module)
+
+
+LINENO_MAPPING = {
+    3: "uint128",
+    4: "uint256",
+    6: "int128",
+    7: "int256",
+    9: "float80",
+    10: "float96",
+    11: "float128",
+    12: "float256",
+    14: "complex160",
+    15: "complex192",
+    16: "complex256",
+    17: "complex512",
+}
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(NO_MYPY, reason="Mypy is not installed")
+def test_extended_precision() -> None:
+    path = os.path.join(MISC_DIR, "extended_precision.py")
+    output_mypy = OUTPUT_MYPY
+    assert path in output_mypy
+
+    for _msg in output_mypy[path]:
+        *_, _lineno, msg_typ, msg = _msg.split(":")
+        lineno = int(_lineno)
+        msg_typ = msg_typ.strip()
+        assert msg_typ in {"error", "note"}
+
+        if LINENO_MAPPING[lineno] in _EXTENDED_PRECISION_LIST:
+            if msg_typ == "error":
+                raise ValueError(f"Unexpected reveal line format: {lineno}")
+            else:
+                marker = FORMAT_DICT[LINENO_MAPPING[lineno]]
+                _test_reveal(path, marker, msg, lineno)
+        else:
+            if msg_typ == "error":
+                marker = "Module has no attribute"
+                _test_fail(path, marker, msg, lineno)