Production hardening: input validation, DRY memo, import efficiency

- Add ShapeMemo.snapshot()/restore() to eliminate duplicated 10-line
  save/restore pattern in _StructChecker and _ArrayLikeChecker
- Validate `casting` parameter in make_array_like_type and
  make_scalar_like_type (reject invalid strings early)
- Validate `byteorder` parameter in DtypeSpec.__post_init__
- Move XLikeNumpy imports under TYPE_CHECKING in jax.py/torch.py
  (23 fewer attribute lookups at runtime per backend import)
- Tighten _check_fixed_dims signature (remove unused list type)
- Add test_coverage_edges.py to numpy tox backend set in conftest.py
- Add tests for all new validation paths + snapshot/restore

Author: acecchini

src/shapix/_array_types.py

@@ -45,6 +45,8 @@
 
 __all__ = ["make_array_type", "make_array_like_type"]
 
+_VALID_CASTINGS = frozenset({"no", "equiv", "safe", "same_kind", "unsafe"})
+
 
 # ---------------------------------------------------------------------------
 # Validator callable (used inside beartype's Is[...])
@@ -90,22 +92,12 @@ def __call__(self, obj: object) -> bool:
       return False
 
     memo = get_memo(_depth=3)
-
-    # Snapshot memo state so we can restore on failure (avoid polluting
-    # the memo with partial bindings from a bad argument).
-    single_snap = memo.single.copy()
-    variadic_snap = memo.variadic.copy()
-    structures_snap = memo.structures.copy()
+    snap = memo.snapshot()
 
     result = check_shape(tuple(shape), self._shape_spec, memo) == ""
 
     if not result:
-      memo.single.clear()
-      memo.single.update(single_snap)
-      memo.variadic.clear()
-      memo.variadic.update(variadic_snap)
-      memo.structures.clear()
-      memo.structures.update(structures_snap)
+      memo.restore(snap)
       self._fail_obj = obj
 
     return result
@@ -242,21 +234,10 @@ def _check(self, obj: object, shape: tuple[int, ...], memo: ShapeMemo) -> bool:
     if not self._check_dtype(obj):
       return False
 
-    # Snapshot memo state so we can restore on failure
-    single_snap = memo.single.copy()
-    variadic_snap = memo.variadic.copy()
-    structures_snap = memo.structures.copy()
-
+    snap = memo.snapshot()
     result = check_shape(shape, self._shape_spec, memo) == ""
-
     if not result:
-      memo.single.clear()
-      memo.single.update(single_snap)
-      memo.variadic.clear()
-      memo.variadic.update(variadic_snap)
-      memo.structures.clear()
-      memo.structures.update(structures_snap)
-
+      memo.restore(snap)
     return result
 
   def _check_dtype(self, obj: object) -> bool:
@@ -350,6 +331,9 @@ def make_array_like_type(
       F32Like = make_array_like_type(FLOAT32, name="F32Like")
       F32Like[N, C, H, W]  # → Annotated[object, Is[...]]
   """
+  if casting not in _VALID_CASTINGS:
+    msg = f"Invalid casting {casting!r}, must be one of {sorted(_VALID_CASTINGS)}"
+    raise ValueError(msg)
   return _ArrayLikeFactory(dtype_spec, casting=casting, name=name)
 
 

src/shapix/_dtypes.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+import typing as tp
 from dataclasses import dataclass, field
 
 __all__ = [
@@ -159,6 +160,18 @@ class DtypeSpec:
   byteorder: str = "any"
   _structured: object = field(default=None, repr=False, compare=False)
 
+  _VALID_BYTEORDERS: tp.ClassVar[frozenset[str]] = frozenset({
+    "any",
+    "little",
+    "big",
+    "native",
+  })
+
+  def __post_init__(self) -> None:
+    if self.byteorder not in self._VALID_BYTEORDERS:
+      msg = f"Invalid byteorder {self.byteorder!r}, must be one of {sorted(self._VALID_BYTEORDERS)}"
+      raise ValueError(msg)
+
   def matches(self, obj: object) -> bool:
     """Return ``True`` if *obj*'s dtype matches this spec."""
     dtype_str = extract_dtype_str(obj)

src/shapix/_memo.py

@@ -35,6 +35,29 @@ class ShapeMemo:
   structures: dict[str, object] = field(default_factory=dict)
   """Tree structure bindings: ``{"T": <TreeSpec>}``."""
 
+  def snapshot(
+    self,
+  ) -> tuple[
+    dict[str, int], dict[str, tuple[bool, tuple[int, ...]]], dict[str, object]
+  ]:
+    """Capture a copy of all current bindings."""
+    return self.single.copy(), self.variadic.copy(), self.structures.copy()
+
+  def restore(
+    self,
+    snap: tuple[
+      dict[str, int], dict[str, tuple[bool, tuple[int, ...]]], dict[str, object]
+    ],
+  ) -> None:
+    """Roll back all bindings to a previous snapshot."""
+    single, variadic, structures = snap
+    self.single.clear()
+    self.single.update(single)
+    self.variadic.clear()
+    self.variadic.update(variadic)
+    self.structures.clear()
+    self.structures.update(structures)
+
 
 _local = threading.local()
 

src/shapix/_shape.py

@@ -194,7 +194,7 @@ def check_shape(
 
 
 def _check_fixed_dims(
-  spec: tuple[DimSpec, ...] | list[DimSpec], shape: tuple[int, ...], memo: ShapeMemo
+  spec: tuple[DimSpec, ...], shape: tuple[int, ...], memo: ShapeMemo
 ) -> str:
   for dim, size in zip(spec, shape):
     err = _check_one(dim, size, memo)

src/shapix/jax.py

@@ -211,36 +211,7 @@ def forward(x: F32[N, C, H, W]) -> BF16[N, C, H, W]: ...
 # Like types (scalar | array | nested sequences — for input validation)
 # ---------------------------------------------------------------------------
 
-from .numpy import (
-  BoolLike as BoolLikeNumpy,
-  C64Like as C64LikeNumpy,
-  C128Like as C128LikeNumpy,
-  ComplexLike as ComplexLikeNumpy,
-  F16Like as F16LikeNumpy,
-  F32Like as F32LikeNumpy,
-  F64Like as F64LikeNumpy,
-  FloatLike as FloatLikeNumpy,
-  InexactLike as InexactLikeNumpy,
-  I8Like as I8LikeNumpy,
-  I16Like as I16LikeNumpy,
-  I32Like as I32LikeNumpy,
-  I64Like as I64LikeNumpy,
-  IntegerLike as IntegerLikeNumpy,
-  IntLike as IntLikeNumpy,
-  NumLike as NumLikeNumpy,
-  RealLike as RealLikeNumpy,
-  ShapedLike as ShapedLikeNumpy,
-  U8Like as U8LikeNumpy,
-  U16Like as U16LikeNumpy,
-  U32Like as U32LikeNumpy,
-  U64Like as U64LikeNumpy,
-  UIntLike as UIntLikeNumpy,
-)
-
-# ---------------------------------------------------------------------------
-# ScalarLike types (re-exported from numpy — no shape, just value validation)
-# ---------------------------------------------------------------------------
-
+# ScalarLike types + factory (re-exported from numpy — no shape, just value)
 from .numpy import BoolScalarLike as BoolScalarLike
 from .numpy import C64ScalarLike as C64ScalarLike
 from .numpy import C128ScalarLike as C128ScalarLike
@@ -268,6 +239,32 @@ def forward(x: F32[N, C, H, W]) -> BF16[N, C, H, W]: ...
 from .numpy import make_scalar_like_type as make_scalar_like_type
 
 if tp.TYPE_CHECKING:
+  from .numpy import (
+    BoolLike as BoolLikeNumpy,
+    C64Like as C64LikeNumpy,
+    C128Like as C128LikeNumpy,
+    ComplexLike as ComplexLikeNumpy,
+    F16Like as F16LikeNumpy,
+    F32Like as F32LikeNumpy,
+    F64Like as F64LikeNumpy,
+    FloatLike as FloatLikeNumpy,
+    InexactLike as InexactLikeNumpy,
+    I8Like as I8LikeNumpy,
+    I16Like as I16LikeNumpy,
+    I32Like as I32LikeNumpy,
+    I64Like as I64LikeNumpy,
+    IntegerLike as IntegerLikeNumpy,
+    IntLike as IntLikeNumpy,
+    NumLike as NumLikeNumpy,
+    RealLike as RealLikeNumpy,
+    ShapedLike as ShapedLikeNumpy,
+    U8Like as U8LikeNumpy,
+    U16Like as U16LikeNumpy,
+    U32Like as U32LikeNumpy,
+    U64Like as U64LikeNumpy,
+    UIntLike as UIntLikeNumpy,
+  )
+
   type BoolLike[*Dims] = Bool[*Dims] | BoolLikeNumpy[*Dims]
 
   type I8Like[*Dims] = I8[*Dims] | I8LikeNumpy[*Dims]

src/shapix/numpy.py

@@ -463,6 +463,12 @@ def make_scalar_like_type(
   type
       An ``Annotated`` type that beartype validates at runtime.
   """
+  from ._array_types import _VALID_CASTINGS
+
+  if casting not in _VALID_CASTINGS:
+    msg = f"Invalid casting {casting!r}, must be one of {sorted(_VALID_CASTINGS)}"
+    raise ValueError(msg)
+
   target = np.dtype(target_dtype)
 
   def _check(value: object) -> bool:

src/shapix/torch.py

@@ -206,36 +206,7 @@ def forward(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 # Like types (scalar | tensor | nested sequences — for input validation)
 # ---------------------------------------------------------------------------
 
-from .numpy import (
-  BoolLike as BoolLikeNumpy,
-  C64Like as C64LikeNumpy,
-  C128Like as C128LikeNumpy,
-  ComplexLike as ComplexLikeNumpy,
-  F16Like as F16LikeNumpy,
-  F32Like as F32LikeNumpy,
-  F64Like as F64LikeNumpy,
-  FloatLike as FloatLikeNumpy,
-  InexactLike as InexactLikeNumpy,
-  I8Like as I8LikeNumpy,
-  I16Like as I16LikeNumpy,
-  I32Like as I32LikeNumpy,
-  I64Like as I64LikeNumpy,
-  IntegerLike as IntegerLikeNumpy,
-  IntLike as IntLikeNumpy,
-  NumLike as NumLikeNumpy,
-  RealLike as RealLikeNumpy,
-  ShapedLike as ShapedLikeNumpy,
-  U8Like as U8LikeNumpy,
-  U16Like as U16LikeNumpy,
-  U32Like as U32LikeNumpy,
-  U64Like as U64LikeNumpy,
-  UIntLike as UIntLikeNumpy,
-)
-
-# ---------------------------------------------------------------------------
-# ScalarLike types (re-exported from numpy — no shape, just value validation)
-# ---------------------------------------------------------------------------
-
+# ScalarLike types + factory (re-exported from numpy — no shape, just value)
 from .numpy import BoolScalarLike as BoolScalarLike
 from .numpy import C64ScalarLike as C64ScalarLike
 from .numpy import C128ScalarLike as C128ScalarLike
@@ -263,6 +234,32 @@ def forward(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 from .numpy import make_scalar_like_type as make_scalar_like_type
 
 if tp.TYPE_CHECKING:
+  from .numpy import (
+    BoolLike as BoolLikeNumpy,
+    C64Like as C64LikeNumpy,
+    C128Like as C128LikeNumpy,
+    ComplexLike as ComplexLikeNumpy,
+    F16Like as F16LikeNumpy,
+    F32Like as F32LikeNumpy,
+    F64Like as F64LikeNumpy,
+    FloatLike as FloatLikeNumpy,
+    InexactLike as InexactLikeNumpy,
+    I8Like as I8LikeNumpy,
+    I16Like as I16LikeNumpy,
+    I32Like as I32LikeNumpy,
+    I64Like as I64LikeNumpy,
+    IntegerLike as IntegerLikeNumpy,
+    IntLike as IntLikeNumpy,
+    NumLike as NumLikeNumpy,
+    RealLike as RealLikeNumpy,
+    ShapedLike as ShapedLikeNumpy,
+    U8Like as U8LikeNumpy,
+    U16Like as U16LikeNumpy,
+    U32Like as U32LikeNumpy,
+    U64Like as U64LikeNumpy,
+    UIntLike as UIntLikeNumpy,
+  )
+
   type BoolLike[*Dims] = Bool[*Dims] | BoolLikeNumpy[*Dims]
 
   type I8Like[*Dims] = I8[*Dims] | I8LikeNumpy[*Dims]

tests/conftest.py

@@ -17,6 +17,7 @@
     "test_memo.py",
     "test_shape.py",
     "test_dimensions.py",
+    "test_coverage_edges.py",
   },
   "jax": {"test_jax.py", "test_dtypes.py"},
   "torch": {"test_torch.py", "test_dtypes.py"},

tests/test_coverage_edges.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import numpy as np
+import pytest
 
 import shapix._memo as memo_mod
 from beartype import BeartypeConf
@@ -288,6 +289,47 @@ def test_struct_checker_fail_obj_replay(self) -> None:
     assert checker(bad) is False  # replay
 
 
+class TestInputValidation:
+  def test_invalid_casting_in_make_array_like_type(self) -> None:
+    from shapix._array_types import make_array_like_type
+
+    with pytest.raises(ValueError, match="Invalid casting"):
+      make_array_like_type(FLOAT32, casting="bogus")
+
+  def test_invalid_casting_in_make_scalar_like_type(self) -> None:
+    from shapix.numpy import make_scalar_like_type
+
+    with pytest.raises(ValueError, match="Invalid casting"):
+      make_scalar_like_type(np.float32, casting="bogus")
+
+  def test_invalid_byteorder_in_dtype_spec(self) -> None:
+    from shapix._dtypes import DtypeSpec
+
+    with pytest.raises(ValueError, match="Invalid byteorder"):
+      DtypeSpec("Bad", frozenset({"float32"}), byteorder="wrong")
+
+  def test_valid_castings_accepted(self) -> None:
+    from shapix._array_types import make_array_like_type
+
+    for casting in ("no", "equiv", "safe", "same_kind", "unsafe"):
+      factory = make_array_like_type(FLOAT32, casting=casting)
+      assert factory is not None
+
+  def test_memo_snapshot_restore(self) -> None:
+    memo = ShapeMemo(single={"N": 5}, variadic={"B": (False, (2, 3))})
+    memo.structures["T"] = "spec"
+    snap = memo.snapshot()
+
+    memo.single["C"] = 10
+    memo.variadic["X"] = (True, (1,))
+    memo.structures["S"] = "other"
+
+    memo.restore(snap)
+    assert memo.single == {"N": 5}
+    assert memo.variadic == {"B": (False, (2, 3))}
+    assert memo.structures == {"T": "spec"}
+
+
 class TestClawWrapper:
   def test_shapix_this_package_delegates_to_beartype(self, monkeypatch) -> None:
     captured: dict[str, object] = {}