summaryrefslogtreecommitdiff
path: root/numpy/typing
diff options
context:
space:
mode:
Diffstat (limited to 'numpy/typing')
-rw-r--r--numpy/typing/__init__.py11
-rw-r--r--numpy/typing/mypy_plugin.py62
-rw-r--r--numpy/typing/tests/data/reveal/ctypeslib.py3
-rw-r--r--numpy/typing/tests/data/reveal/ndarray_misc.py4
-rw-r--r--numpy/typing/tests/test_typing.py9
5 files changed, 67 insertions, 22 deletions
diff --git a/numpy/typing/__init__.py b/numpy/typing/__init__.py
index 04d34f0c7..252123a19 100644
--- a/numpy/typing/__init__.py
+++ b/numpy/typing/__init__.py
@@ -23,15 +23,18 @@ Mypy plugin
-----------
A mypy_ plugin is distributed in `numpy.typing` for managing a number of
-platform-specific annotations. Its function can be split into to parts:
+platform-specific annotations. Its functionality can be split into three
+distinct 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`.
+ comprehensive overview of the affected classes. Without the plugin the
+ precision of all relevant classes will be inferred as `~typing.Any`.
+* Assigning the (platform-dependent) precision of `~numpy.ctypeslib.c_intp`.
+ Without the plugin aforementioned type will default to `ctypes.c_int64`.
* Removing all extended-precision `~numpy.number` subclasses that are unavailable
- for the platform in question. Most notable this includes the likes of
+ for the platform in question. Most notably, 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.
diff --git a/numpy/typing/mypy_plugin.py b/numpy/typing/mypy_plugin.py
index 100e0d957..74dcd7a85 100644
--- a/numpy/typing/mypy_plugin.py
+++ b/numpy/typing/mypy_plugin.py
@@ -61,6 +61,19 @@ def _get_extended_precision_list() -> t.List[str]:
return [i.__name__ for i in extended_types if i.__name__ in extended_names]
+def _get_c_intp_name() -> str:
+ # Adapted from `np.core._internal._getintp_ctype`
+ char = np.dtype('p').char
+ if char == 'i':
+ return "c_int"
+ elif char == 'l':
+ return "c_long"
+ elif char == 'q':
+ return "c_longlong"
+ else:
+ return "c_long"
+
+
#: A dictionary mapping type-aliases in `numpy.typing._nbit` to
#: concrete `numpy.typing.NBitBase` subclasses.
_PRECISION_DICT: t.Final = _get_precision_dict()
@@ -68,6 +81,9 @@ _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()
+#: The name of the ctypes quivalent of `np.intp`
+_C_INTP: t.Final = _get_c_intp_name()
+
def _hook(ctx: AnalyzeTypeContext) -> Type:
"""Replace a type-alias with a concrete ``NBitBase`` subclass."""
@@ -87,8 +103,23 @@ if t.TYPE_CHECKING or MYPY_EX is None:
raise ValueError("Failed to identify a `ImportFrom` instance "
f"with the following id: {id!r}")
+ def _override_imports(
+ file: MypyFile,
+ module: str,
+ imports: t.List[t.Tuple[str, t.Optional[str]]],
+ ) -> None:
+ """Override the first `module`-based import with new `imports`."""
+ # Construct a new `from module import y` statement
+ import_obj = ImportFrom(module, 0, names=imports)
+ import_obj.is_top_level = True
+
+ # Replace the first `module`-based import statement with `import_obj`
+ for lst in [file.defs, file.imports]: # type: t.List[Statement]
+ i = _index(lst, module)
+ lst[i] = import_obj
+
class _NumpyPlugin(Plugin):
- """A plugin for assigning platform-specific `numpy.number` precisions."""
+ """A mypy plugin for handling versus numpy-specific typing tasks."""
def get_type_analyze_hook(self, fullname: str) -> None | _HookFunc:
"""Set the precision of platform-specific `numpy.number` subclasses.
@@ -100,25 +131,26 @@ if t.TYPE_CHECKING or MYPY_EX is None:
return None
def get_additional_deps(self, file: MypyFile) -> t.List[t.Tuple[int, str, int]]:
- """Import platform-specific extended-precision `numpy.number` subclasses.
+ """Handle all import-based overrides.
+
+ * Import platform-specific extended-precision `numpy.number`
+ subclasses (*e.g.* `numpy.float96`, `numpy.float128` and
+ `numpy.complex256`).
+ * Import the appropriate `ctypes` equivalent to `numpy.intp`.
- 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],
+ _override_imports(
+ file, "numpy.typing._extended_precision",
+ imports=[(v, v) for v in _EXTENDED_PRECISION_LIST],
+ )
+ elif file.fullname == "numpy.ctypeslib":
+ _override_imports(
+ file, "ctypes",
+ imports=[(_C_INTP, "_c_intp")],
)
- 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]:
diff --git a/numpy/typing/tests/data/reveal/ctypeslib.py b/numpy/typing/tests/data/reveal/ctypeslib.py
new file mode 100644
index 000000000..0c32d70ed
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/ctypeslib.py
@@ -0,0 +1,3 @@
+import numpy as np
+
+reveal_type(np.ctypeslib.c_intp()) # E: {c_intp}
diff --git a/numpy/typing/tests/data/reveal/ndarray_misc.py b/numpy/typing/tests/data/reveal/ndarray_misc.py
index ea01b7aa4..2e198eb6f 100644
--- a/numpy/typing/tests/data/reveal/ndarray_misc.py
+++ b/numpy/typing/tests/data/reveal/ndarray_misc.py
@@ -23,8 +23,8 @@ AR_U: np.ndarray[Any, np.dtype[np.str_]]
ctypes_obj = AR_f8.ctypes
reveal_type(ctypes_obj.data) # E: int
-reveal_type(ctypes_obj.shape) # E: ctypes.Array[ctypes.c_int64]
-reveal_type(ctypes_obj.strides) # E: ctypes.Array[ctypes.c_int64]
+reveal_type(ctypes_obj.shape) # E: ctypes.Array[{c_intp}]
+reveal_type(ctypes_obj.strides) # E: ctypes.Array[{c_intp}]
reveal_type(ctypes_obj._as_parameter_) # E: ctypes.c_void_p
reveal_type(ctypes_obj.data_as(ct.c_void_p)) # E: ctypes.c_void_p
diff --git a/numpy/typing/tests/test_typing.py b/numpy/typing/tests/test_typing.py
index be08c1359..35558c880 100644
--- a/numpy/typing/tests/test_typing.py
+++ b/numpy/typing/tests/test_typing.py
@@ -8,7 +8,11 @@ from typing import Optional, IO, Dict, List
import pytest
import numpy as np
-from numpy.typing.mypy_plugin import _PRECISION_DICT, _EXTENDED_PRECISION_LIST
+from numpy.typing.mypy_plugin import (
+ _PRECISION_DICT,
+ _EXTENDED_PRECISION_LIST,
+ _C_INTP,
+)
try:
from mypy import api
@@ -219,6 +223,9 @@ def _construct_format_dict():
# numpy.typing
"_NBitInt": dct['_NBitInt'],
+
+ # numpy.ctypeslib
+ "c_intp": f"ctypes.{_C_INTP}"
}