Skip to content

Commit 9ca8076

Browse files
BvB93Bas van Beek
authored andcommitted
ENH: Add number.__class_getitem__
1 parent 07124b5 commit 9ca8076

File tree

4 files changed

+129
-0
lines changed

4 files changed

+129
-0
lines changed

numpy/__init__.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3058,6 +3058,8 @@ class number(generic, Generic[_NBit1]): # type: ignore
30583058
def real(self: _ArraySelf) -> _ArraySelf: ...
30593059
@property
30603060
def imag(self: _ArraySelf) -> _ArraySelf: ...
3061+
if sys.version_info >= (3, 9):
3062+
def __class_getitem__(self, item: Any) -> GenericAlias: ...
30613063
def __int__(self) -> int: ...
30623064
def __float__(self) -> float: ...
30633065
def __complex__(self) -> complex: ...

numpy/core/_add_newdocs.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6500,6 +6500,37 @@ def refer_to_array_attribute(attr, method=True):
65006500
add_newdoc('numpy.core.numerictypes', 'generic',
65016501
refer_to_array_attribute('view'))
65026502

6503+
if sys.version_info >= (3, 9):
6504+
add_newdoc('numpy.core.numerictypes', 'number', ('__class_getitem__',
6505+
"""
6506+
__class_getitem__(item, /)
6507+
6508+
Return a parametrized wrapper around the `~numpy.number` type.
6509+
6510+
.. versionadded:: 1.22
6511+
6512+
Returns
6513+
-------
6514+
alias : types.GenericAlias
6515+
A parametrized `~numpy.number` type.
6516+
6517+
Examples
6518+
--------
6519+
>>> from typing import Any
6520+
>>> import numpy as np
6521+
6522+
>>> np.signedinteger[Any]
6523+
numpy.signedinteger[typing.Any]
6524+
6525+
Note
6526+
----
6527+
This method is only available for python 3.9 and later.
6528+
6529+
See Also
6530+
--------
6531+
:pep:`585` : Type hinting generics in standard collections.
6532+
6533+
"""))
65036534

65046535
##############################################################################
65056536
#

numpy/core/src/multiarray/scalartypes.c.src

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1805,6 +1805,21 @@ gentype_setflags(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args),
18051805
Py_RETURN_NONE;
18061806
}
18071807

1808+
/*
1809+
* Use for concrete np.number subclasses, making them act as if they
1810+
* were subtyped from e.g. np.signedinteger[object], thus lacking any
1811+
* free subscription parameters. Requires python >= 3.9.
1812+
*/
1813+
#ifdef Py_GENERICALIASOBJECT_H
1814+
static PyObject *
1815+
numbertype_class_getitem(PyObject *cls, PyObject *args)
1816+
{
1817+
return PyErr_Format(PyExc_TypeError,
1818+
"There are no type variables left in %s",
1819+
((PyTypeObject *)cls)->tp_name);
1820+
}
1821+
#endif
1822+
18081823
/*
18091824
* casting complex numbers (that don't inherit from Python complex)
18101825
* to Python complex
@@ -2188,13 +2203,29 @@ static PyGetSetDef inttype_getsets[] = {
21882203
{NULL, NULL, NULL, NULL, NULL}
21892204
};
21902205

2206+
static PyMethodDef numbertype_methods[] = {
2207+
/* for typing; requires python >= 3.9 */
2208+
#ifdef Py_GENERICALIASOBJECT_H
2209+
{"__class_getitem__",
2210+
(PyCFunction)Py_GenericAlias,
2211+
METH_CLASS | METH_O, NULL},
2212+
#endif
2213+
{NULL, NULL, 0, NULL} /* sentinel */
2214+
};
2215+
21912216
/**begin repeat
21922217
* #name = cfloat,clongdouble#
21932218
*/
21942219
static PyMethodDef @name@type_methods[] = {
21952220
{"__complex__",
21962221
(PyCFunction)@name@_complex,
21972222
METH_VARARGS | METH_KEYWORDS, NULL},
2223+
/* for typing; requires python >= 3.9 */
2224+
#ifdef Py_GENERICALIASOBJECT_H
2225+
{"__class_getitem__",
2226+
(PyCFunction)numbertype_class_getitem,
2227+
METH_CLASS | METH_O, NULL},
2228+
#endif
21982229
{NULL, NULL, 0, NULL}
21992230
};
22002231
/**end repeat**/
@@ -2232,6 +2263,27 @@ static PyMethodDef @name@type_methods[] = {
22322263
{"is_integer",
22332264
(PyCFunction)@name@_is_integer,
22342265
METH_NOARGS, NULL},
2266+
/* for typing; requires python >= 3.9 */
2267+
#ifdef Py_GENERICALIASOBJECT_H
2268+
{"__class_getitem__",
2269+
(PyCFunction)numbertype_class_getitem,
2270+
METH_CLASS | METH_O, NULL},
2271+
#endif
2272+
{NULL, NULL, 0, NULL}
2273+
};
2274+
/**end repeat**/
2275+
2276+
/**begin repeat
2277+
* #name = byte, short, int, long, longlong, ubyte, ushort,
2278+
* uint, ulong, ulonglong, timedelta, cdouble#
2279+
*/
2280+
static PyMethodDef @name@type_methods[] = {
2281+
/* for typing; requires python >= 3.9 */
2282+
#ifdef Py_GENERICALIASOBJECT_H
2283+
{"__class_getitem__",
2284+
(PyCFunction)numbertype_class_getitem,
2285+
METH_CLASS | METH_O, NULL},
2286+
#endif
22352287
{NULL, NULL, 0, NULL}
22362288
};
22372289
/**end repeat**/
@@ -3951,6 +4003,8 @@ initialize_numeric_types(void)
39514003

39524004
PyIntegerArrType_Type.tp_getset = inttype_getsets;
39534005

4006+
PyNumberArrType_Type.tp_methods = numbertype_methods;
4007+
39544008
/**begin repeat
39554009
* #NAME= Number, Integer, SignedInteger, UnsignedInteger, Inexact,
39564010
* Floating, ComplexFloating, Flexible, Character#
@@ -4016,6 +4070,17 @@ initialize_numeric_types(void)
40164070

40174071
/**end repeat**/
40184072

4073+
/**begin repeat
4074+
* #name = byte, short, int, long, longlong, ubyte, ushort,
4075+
* uint, ulong, ulonglong, timedelta, cdouble#
4076+
* #Name = Byte, Short, Int, Long, LongLong, UByte, UShort,
4077+
* UInt, ULong, ULongLong, Timedelta, CDouble#
4078+
*/
4079+
4080+
Py@Name@ArrType_Type.tp_methods = @name@type_methods;
4081+
4082+
/**end repeat**/
4083+
40194084
/* We won't be inheriting from Python Int type. */
40204085
PyIntArrType_Type.tp_hash = int_arrtype_hash;
40214086

numpy/core/tests/test_scalar_methods.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
"""
22
Test the scalar constructors, which also do type-coercion
33
"""
4+
import sys
45
import fractions
56
import platform
7+
import types
8+
from typing import Any, Type
69

710
import pytest
811
import numpy as np
@@ -128,3 +131,31 @@ def test_false(self, code: str) -> None:
128131
if value == 0:
129132
continue
130133
assert not value.is_integer()
134+
135+
136+
@pytest.mark.skipif(sys.version_info < (3, 9), reason="Requires python 3.9")
137+
class TestClassGetItem:
138+
@pytest.mark.parametrize("cls", [
139+
np.number,
140+
np.integer,
141+
np.inexact,
142+
np.unsignedinteger,
143+
np.signedinteger,
144+
np.floating,
145+
np.complexfloating,
146+
])
147+
def test_abc(self, cls: Type[np.number]) -> None:
148+
alias = cls[Any]
149+
assert isinstance(alias, types.GenericAlias)
150+
assert alias.__origin__ is cls
151+
152+
@pytest.mark.parametrize("cls", [np.generic, np.flexible, np.character])
153+
def test_abc_non_numeric(self, cls: Type[np.generic]) -> None:
154+
with pytest.raises(TypeError):
155+
cls[Any]
156+
157+
@pytest.mark.parametrize("code", np.typecodes["All"])
158+
def test_concrete(self, code: str) -> None:
159+
cls = np.dtype(code).type
160+
with pytest.raises(TypeError):
161+
cls[Any]

0 commit comments

Comments
 (0)