ENH: Added a mypy plugin for handling platform-specific `np.number` precisions

7 files changed, 229 insertions, 37 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 3165a6319..6df10d7a7 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -21,10 +21,24 @@ from numpy.typing import (
     _SupportsDType,
     _VoidDTypeLike,
     NBitBase,
+    _256Bit,
+    _128Bit,
+    _96Bit,
+    _80Bit,
     _64Bit,
     _32Bit,
     _16Bit,
     _8Bit,
+    _NBitByte,
+    _NBitShort,
+    _NBitIntC,
+    _NBitIntP,
+    _NBitInt,
+    _NBitLongLong,
+    _NBitHalf,
+    _NBitSingle,
+    _NBitDouble,
+    _NBitLongDouble,
 )
 from numpy.typing._callable import (
     _BoolOp,
@@ -311,29 +325,22 @@ broadcast_to: Any
 busday_count: Any
 busday_offset: Any
 busdaycalendar: Any
-byte: Any
 byte_bounds: Any
 bytes0: Any
 c_: Any
 can_cast: Any
 cast: Any
-cdouble: Any
-cfloat: Any
 chararray: Any
-clongdouble: Any
-clongfloat: Any
 column_stack: Any
 common_type: Any
 compare_chararrays: Any
 complex256: Any
-complex_: Any
 concatenate: Any
 conj: Any
 copy: Any
 copyto: Any
 corrcoef: Any
 cov: Any
-csingle: Any
 cumproduct: Any
 datetime_as_string: Any
 datetime_data: Any
@@ -349,7 +356,6 @@ digitize: Any
 disp: Any
 divide: Any
 dot: Any
-double: Any
 dsplit: Any
 dstack: Any
 ediff1d: Any
@@ -365,7 +371,6 @@ flip: Any
 fliplr: Any
 flipud: Any
 float128: Any
-float_: Any
 format_parser: Any
 frombuffer: Any
 fromfile: Any
@@ -377,7 +382,6 @@ genfromtxt: Any
 get_include: Any
 geterrobj: Any
 gradient: Any
-half: Any
 hamming: Any
 hanning: Any
 histogram: Any
@@ -393,12 +397,8 @@ index_exp: Any
 info: Any
 inner: Any
 insert: Any
-int0: Any
-int_: Any
-intc: Any
 interp: Any
 intersect1d: Any
-intp: Any
 is_busday: Any
 iscomplex: Any
 iscomplexobj: Any
@@ -415,10 +415,6 @@ lexsort: Any
 load: Any
 loads: Any
 loadtxt: Any
-longcomplex: Any
-longdouble: Any
-longfloat: Any
-longlong: Any
 lookfor: Any
 mafromtxt: Any
 mask_indices: Any
@@ -504,11 +500,8 @@ setdiff1d: Any
 seterrobj: Any
 setxor1d: Any
 shares_memory: Any
-short: Any
 show_config: Any
 sinc: Any
-single: Any
-singlecomplex: Any
 sort_complex: Any
 source: Any
 split: Any
@@ -527,18 +520,11 @@ triu_indices_from: Any
 typeDict: Any
 typecodes: Any
 typename: Any
-ubyte: Any
-uint: Any
-uint0: Any
-uintc: Any
-uintp: Any
-ulonglong: Any
 union1d: Any
 unique: Any
 unpackbits: Any
 unravel_index: Any
 unwrap: Any
-ushort: Any
 vander: Any
 vdot: Any
 vectorize: Any
@@ -1853,6 +1839,14 @@ int16 = signedinteger[_16Bit]
 int32 = signedinteger[_32Bit]
 int64 = signedinteger[_64Bit]
 
+byte = signedinteger[_NBitByte]
+short = signedinteger[_NBitShort]
+intc = signedinteger[_NBitIntC]
+intp = signedinteger[_NBitIntP]
+int0 = signedinteger[_NBitIntP]
+int_ = signedinteger[_NBitInt]
+longlong = signedinteger[_NBitLongLong]
+
 class timedelta64(generic):
     def __init__(
         self,
@@ -1917,6 +1911,14 @@ uint16 = unsignedinteger[_16Bit]
 uint32 = unsignedinteger[_32Bit]
 uint64 = unsignedinteger[_64Bit]
 
+ubyte = unsignedinteger[_NBitByte]
+ushort = unsignedinteger[_NBitShort]
+uintc = unsignedinteger[_NBitIntC]
+uintp = unsignedinteger[_NBitIntP]
+uint0 = unsignedinteger[_NBitIntP]
+uint = unsignedinteger[_NBitInt]
+ulonglong = unsignedinteger[_NBitLongLong]
+
 class inexact(number[_NBit_co]): ...  # type: ignore
 
 _IntType = TypeVar("_IntType", bound=integer)
@@ -1945,6 +1947,13 @@ float16 = floating[_16Bit]
 float32 = floating[_32Bit]
 float64 = floating[_64Bit]
 
+half = floating[_NBitHalf]
+single = floating[_NBitSingle]
+double = floating[_NBitDouble]
+float_ = floating[_NBitDouble]
+longdouble = floating[_NBitLongDouble]
+longfloat = floating[_NBitLongDouble]
+
 # The main reason for `complexfloating` having two typevars is cosmetic.
 # It is used to clarify why `complex128`s precision is `_64Bit`, the latter
 # describing the two 64 bit floats representing its real and imaginary component
@@ -1972,6 +1981,15 @@ class complexfloating(inexact[_NBit_co], Generic[_NBit_co, _NBit_co2]):
 complex64 = complexfloating[_32Bit, _32Bit]
 complex128 = complexfloating[_64Bit, _64Bit]
 
+csingle = complexfloating[_NBitSingle, _NBitSingle]
+singlecomplex = complexfloating[_NBitSingle, _NBitSingle]
+cdouble = complexfloating[_NBitDouble, _NBitDouble]
+complex_ = complexfloating[_NBitDouble, _NBitDouble]
+cfloat = complexfloating[_NBitDouble, _NBitDouble]
+clongdouble = complexfloating[_NBitLongDouble, _NBitLongDouble]
+clongfloat = complexfloating[_NBitLongDouble, _NBitLongDouble]
+longcomplex = complexfloating[_NBitLongDouble, _NBitLongDouble]
+
 class flexible(generic): ...  # type: ignore
 
 class void(flexible):
@@ -2007,15 +2025,10 @@ class str_(character, str):
         self, __value: bytes, encoding: str = ..., errors: str = ...
     ) -> None: ...
 
-unicode_ = str0 = str_
+unicode_ = str_
+str0 = str_
 
-# TODO(alan): Platform dependent types
-# longcomplex, longdouble, longfloat
-# bytes, short, intc, intp, longlong
-# half, single, double, longdouble
-# uint_, int_, float_, complex_
-# float128, complex256
-# float96
+# TODO: Platform dependent types: float128, complex256, float96
 
 def array(
     object: object,
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
index d9d9557bf..694046609 100644
--- a/numpy/typing/__init__.py
+++ b/numpy/typing/__init__.py
@@ -19,6 +19,28 @@ the two below:
 
 .. _typing-extensions: https://pypi.org/project/typing-extensions/
 
+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`.
+
+To enable the plugin, one must add it to their mypy `configuration file`_:
+
+.. code-block:: ini
+
+    [mypy]
+    plugins = numpy.typing.mypy_plugin
+
+.. _mypy: http://mypy-lang.org/
+.. _configuration file: https://mypy.readthedocs.io/en/stable/config_file.html
+
 Differences from the runtime NumPy API
 --------------------------------------
 
@@ -207,6 +229,18 @@ class _8Bit(_16Bit): ...  # type: ignore[misc]
 # Clean up the namespace
 del TYPE_CHECKING, final, List
 
+from ._nbit import (
+    _NBitByte,
+    _NBitShort,
+    _NBitIntC,
+    _NBitIntP,
+    _NBitInt,
+    _NBitLongLong,
+    _NBitHalf,
+    _NBitSingle,
+    _NBitDouble,
+    _NBitLongDouble,
+)
 from ._scalars import (
     _CharLike,
     _BoolLike,
diff --git a/numpy/typing/_nbit.py b/numpy/typing/_nbit.py
new file mode 100644
index 000000000..b8d35db4f
--- /dev/null
+++ b/numpy/typing/_nbit.py
@@ -0,0 +1,16 @@
+"""A module with the precisions of platform-specific `~numpy.number`s."""
+
+from typing import Any
+
+# To-be replaced with a `npt.NBitBase` subclass by numpy's mypy plugin
+_NBitByte = Any
+_NBitShort = Any
+_NBitIntC = Any
+_NBitIntP = Any
+_NBitInt = Any
+_NBitLongLong = Any
+
+_NBitHalf = Any
+_NBitSingle = Any
+_NBitDouble = Any
+_NBitLongDouble = Any
diff --git a/numpy/typing/mypy_plugin.py b/numpy/typing/mypy_plugin.py
new file mode 100644
index 000000000..023ab5fd1
--- /dev/null
+++ b/numpy/typing/mypy_plugin.py
@@ -0,0 +1,58 @@
+"""A module containing `numpy`-specific plugins for mypy."""
+
+import typing as t
+
+import numpy as np
+
+import mypy.types
+from mypy.types import Type
+from mypy.plugin import Plugin, AnalyzeTypeContext
+
+HookFunc = t.Callable[[AnalyzeTypeContext], Type]
+
+
+def _get_precision_dict() -> t.Dict[str, str]:
+    names = [
+        ("_NBitByte", np.byte),
+        ("_NBitShort", np.short),
+        ("_NBitIntC", np.intc),
+        ("_NBitIntP", np.intp),
+        ("_NBitInt", np.int_),
+        ("_NBitLongLong", np.longlong),
+
+        ("_NBitHalf", np.half),
+        ("_NBitSingle", np.single),
+        ("_NBitDouble", np.double),
+        ("_NBitLongDouble", np.longdouble),
+    ]
+    ret = {}
+    for name, typ in names:
+        n: int = 8 * typ().dtype.alignment
+        ret[f'numpy.typing._nbit.{name}'] = f"numpy._{n}Bit"
+    return ret
+
+
+#: A dictionary mapping type-aliases in `numpy.typing._nbit` to
+#: concrete `numpy.typing.NBitBase` subclasses.
+_PRECISION_DICT = _get_precision_dict()
+
+
+def _hook(ctx: AnalyzeTypeContext) -> Type:
+    """Replace a type-alias with a concrete ``NBitBase`` subclass."""
+    typ, _, api = ctx
+    name = typ.name.split(".")[-1]
+    name_new = _PRECISION_DICT[f"numpy.typing._nbit.{name}"]
+    return api.named_type(name_new)
+
+
+class _NumpyPlugin(Plugin):
+    """A plugin for assigning platform-specific `numpy.number` precisions."""
+
+    def get_type_analyze_hook(self, fullname: str) -> t.Optional[HookFunc]:
+        if fullname in _PRECISION_DICT:
+            return _hook
+        return None
+
+
+def plugin(version: str) -> t.Type[_NumpyPlugin]:
+    return _NumpyPlugin
diff --git a/numpy/typing/tests/data/mypy.ini b/numpy/typing/tests/data/mypy.ini
index 91d93588a..35cfbec89 100644
--- a/numpy/typing/tests/data/mypy.ini
+++ b/numpy/typing/tests/data/mypy.ini
@@ -1,5 +1,5 @@
 [mypy]
-mypy_path = ../../..
+plugins = numpy.typing.mypy_plugin
 
 [mypy-numpy]
 ignore_errors = True
diff --git a/numpy/typing/tests/data/pass/scalars.py b/numpy/typing/tests/data/pass/scalars.py
index b7f7880e4..2f2643e8e 100644
--- a/numpy/typing/tests/data/pass/scalars.py
+++ b/numpy/typing/tests/data/pass/scalars.py
@@ -163,3 +163,38 @@ c16.strides
 c16.squeeze()
 c16.byteswap()
 c16.transpose()
+
+# Aliases
+np.str0()
+
+np.byte()
+np.short()
+np.intc()
+np.intp()
+np.int0()
+np.int_()
+np.longlong()
+
+np.ubyte()
+np.ushort()
+np.uintc()
+np.uintp()
+np.uint0()
+np.uint()
+np.ulonglong()
+
+np.half()
+np.single()
+np.double()
+np.float_()
+np.longdouble()
+np.longfloat()
+
+np.csingle()
+np.singlecomplex()
+np.cdouble()
+np.complex_()
+np.cfloat()
+np.clongdouble()
+np.clongfloat()
+np.longcomplex()
diff --git a/numpy/typing/tests/data/reveal/scalars.py b/numpy/typing/tests/data/reveal/scalars.py
index e887e302d..841b1473d 100644
--- a/numpy/typing/tests/data/reveal/scalars.py
+++ b/numpy/typing/tests/data/reveal/scalars.py
@@ -26,3 +26,39 @@ reveal_type(np.complex128().imag)  # E: numpy.floating[numpy.typing._64Bit]
 
 reveal_type(np.unicode_('foo'))  # E: numpy.str_
 reveal_type(np.str0('foo'))  # E: numpy.str_
+
+# Aliases
+reveal_type(np.unicode_())  # E: numpy.str_
+reveal_type(np.str0())  # E: numpy.str_
+
+reveal_type(np.byte())  # E: numpy.signedinteger[numpy.typing._
+reveal_type(np.short())  # E: numpy.signedinteger[numpy.typing._
+reveal_type(np.intc())  # E: numpy.signedinteger[numpy.typing._
+reveal_type(np.intp())  # E: numpy.signedinteger[numpy.typing._
+reveal_type(np.int0())  # E: numpy.signedinteger[numpy.typing._
+reveal_type(np.int_())  # E: numpy.signedinteger[numpy.typing._
+reveal_type(np.longlong())  # E: numpy.signedinteger[numpy.typing._
+
+reveal_type(np.ubyte())  # E: numpy.unsignedinteger[numpy.typing._
+reveal_type(np.ushort())  # E: numpy.unsignedinteger[numpy.typing._
+reveal_type(np.uintc())  # E: numpy.unsignedinteger[numpy.typing._
+reveal_type(np.uintp())  # E: numpy.unsignedinteger[numpy.typing._
+reveal_type(np.uint0())  # E: numpy.unsignedinteger[numpy.typing._
+reveal_type(np.uint())  # E: numpy.unsignedinteger[numpy.typing._
+reveal_type(np.ulonglong())  # E: numpy.unsignedinteger[numpy.typing._
+
+reveal_type(np.half())  # E: numpy.floating[numpy.typing._
+reveal_type(np.single())  # E: numpy.floating[numpy.typing._
+reveal_type(np.double())  # E: numpy.floating[numpy.typing._
+reveal_type(np.float_())  # E: numpy.floating[numpy.typing._
+reveal_type(np.longdouble())  # E: numpy.floating[numpy.typing._
+reveal_type(np.longfloat())  # E: numpy.floating[numpy.typing._
+
+reveal_type(np.csingle())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.singlecomplex())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.cdouble())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.complex_())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.cfloat())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.clongdouble())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.clongfloat())  # E: numpy.complexfloating[numpy.typing._
+reveal_type(np.longcomplex())  # E: numpy.complexfloating[numpy.typing._