Compatibility with Python 3.10 + with generic array types

jaxtyping/_array_types.py

@@ -25,7 +25,16 @@
 import types
 import typing
 from dataclasses import dataclass
-from typing import Any, Literal, NoReturn, Optional, TypeVar, Union
+from typing import (
+ Any,
+ get_args,
+ get_origin,
+ Literal,
+ NoReturn,
+ Optional,
+ TypeVar,
+ Union,
+)
 
 
 # Bit of a hack, but jaxtyping provides nicer error messages than typeguard. This means
@@ -358,7 +367,7 @@ class for `Float32[Array, "foo"]`.
 
 _not_made = object()
 
-_union_types = [typing.Union]
+_union_types = [Union]
 if sys.version_info >= (3, 10):
  _union_types.append(types.UnionType)
 
@@ -517,6 +526,9 @@ def _make_array_cached(array_type, dim_str, dtypes, name):
  # Allow Python built-in numeric types.
  # TODO: do something more generic than this? Should we _make all types
  # that have `shape` and `dtype` attributes or something?
+ array_origin = get_origin(array_type)
+ if array_origin is not None:
+ array_type = array_origin
  if array_type is bool:
  if _check_scalar("bool", dtypes, dims):
  return array_type
@@ -547,7 +559,7 @@ def _make_array_cached(array_type, dim_str, dtypes, name):
  return array_type
  else:
  return _not_made
- if issubclass(array_type, AbstractArray):
+ if array_type is not Any and issubclass(array_type, AbstractArray):
  if dtypes is _any_dtype:
  dtypes = array_type.dtypes
  elif array_type.dtypes is not _any_dtype:
@@ -588,11 +600,15 @@ def _make_array(*args, **kwargs):
 
  if type(out) is tuple:
  array_type, name, dtypes, dims, index_variadic, dim_str = out
- metaclass = _make_metaclass(type(array_type))
+ metaclass = (
+ _make_metaclass(type)
+ if array_type is Any
+ else _make_metaclass(type(array_type))
+ )
 
  out = metaclass(
  name,
- (array_type, AbstractArray),
+ (AbstractArray,) if array_type is Any else (array_type, AbstractArray),
  dict(
  array_type=array_type,
  dtypes=dtypes,
@@ -629,14 +645,18 @@ def __getitem__(cls, item: tuple[Any, str]):
  if isinstance(array_type, TypeVar):
  bound = array_type.__bound__
  if bound is None:
- array_type = Any
+ constraints = array_type.__constraints__
+ if constraints == ():
+ array_type = Any
+ else:
+ array_type = Union[constraints]
  else:
  array_type = bound
  del item
- if typing.get_origin(array_type) in _union_types:
+ if get_origin(array_type) in _union_types:
  out = [
  _make_array(x, dim_str, cls.dtypes, cls.__name__)
- for x in typing.get_args(array_type)
+ for x in get_args(array_type)
  ]
  out = tuple(x for x in out if x is not _not_made)
  if len(out) == 0:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.32"
+version = "0.2.33"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"