is_scalar_type: rule out int subclasses

Author: inducer

arraycontext/container/dataclass.py

@@ -61,6 +61,7 @@
     ArrayContainer,
     ArrayOrContainer,
     ArrayOrContainerOrScalar,
+    is_scalar_type,
 )
 
 
@@ -95,15 +96,6 @@ def _is_array_or_container_type(tp: type | GenericAlias | UnionType, /) -> bool:
     return tp is Array or is_array_container_type(tp)
 
 
-def is_scalar_type(tp: object, /) -> bool:
-    if not isinstance(tp, type):
-        tp = get_origin(tp)
-    if not isinstance(tp, type):
-        return False
-
-    return issubclass(tp, (np.generic, int, float, complex))
-
-
 def dataclass_array_container(cls: type[T]) -> type[T]:
     """A class decorator that makes the class to which it is applied an
     :class:`ArrayContainer` by registering appropriate implementations of

arraycontext/typing.py

@@ -81,6 +81,7 @@
     TypeAlias,
     TypeVar,
     cast,
+    get_origin,
     overload,
 )
 
@@ -94,6 +95,8 @@
 if TYPE_CHECKING:
     from numpy.typing import DTypeLike
 
+    from pymbolic.typing import Integer
+
 
 # deprecated, use ScalarLike instead
 Scalar: TypeAlias = _Scalar
@@ -237,7 +240,7 @@ def __rpow__(self, other: ArrayOrScalar | Self) -> Self: ...
 
 
 ArrayT = TypeVar("ArrayT", bound=Array)
-ArrayOrScalar: TypeAlias = Array | _Scalar
+ArrayOrScalar: TypeAlias = Array | ScalarLike
 ArrayOrScalarT = TypeVar("ArrayOrScalarT", bound=ArrayOrScalar)
 ArrayOrContainer: TypeAlias = Array | ArrayContainer
 ArrayOrArithContainer: TypeAlias = Array | ArithArrayContainer
@@ -261,6 +264,22 @@ def __rpow__(self, other: ArrayOrScalar | Self) -> Self: ...
 NumpyOrContainerOrScalar: TypeAlias = "np.ndarray | ArrayContainer | ScalarLike"
 
 
+def is_scalar_type(tp: object, /) -> bool:
+    if not isinstance(tp, type):
+        tp = get_origin(tp)
+    if not isinstance(tp, type):
+        return False
+    if tp is int or tp is bool:
+        # int has loads of undesirable subclasses: enums, ...
+        # We're not going to tolerate them.
+        #
+        # bool has to be OK because arraycontext is expected to handle
+        # arrays of bools.
+        return True
+
+    return issubclass(tp, (np.generic, float, complex))
+
+
 def is_scalar_like(x: object, /) -> TypeIs[Scalar]:
     return np.isscalar(x)