Refine runtime value shape expression API

Author: acecchini

README.md

@@ -159,7 +159,7 @@ and numeric literals. Attribute access and function calls are rejected.
 
 ### Runtime value dimensions
 
-Use `Value["expr"]` when a shape depends on a runtime parameter or `self`
+Use `Value("expr")` when a shape depends on a runtime parameter or `self`
 attribute rather than a previously bound dimension:
 
 ```python
@@ -168,19 +168,22 @@ from shapix import Value
 from shapix.numpy import F32
 import numpy as np
 
+Size = Value("size")
+WidthPlus3 = Value("self.width + 3")
+
 @beartype
-def full(size: int) -> F32[Value["size"]]:
+def full(size: int) -> F32[Size]:
     return np.full((size,), 1.0, dtype=np.float32)
 
 class SomeClass:
     width = 5
 
     @beartype
-    def full(self) -> F32[Value["self.width + 3"]]:
+    def full(self) -> F32[WidthPlus3]:
         return np.full((self.width + 3,), 1.0, dtype=np.float32)
 ```
 
-`Value[...]` uses a restricted arithmetic grammar as well. It allows names,
+`Value(...)` uses a restricted arithmetic grammar as well. It allows names,
 attribute access, numeric literals, and arithmetic operators, but rejects calls,
 indexing, and other arbitrary Python expressions.
 

src/shapix/__init__.py

@@ -23,7 +23,7 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
     ``B``, ``N``, ``P``, ``L``, ``C``, ``D``, ``K``, ``H``, ``W`` — named dimensions.
     ``__`` — anonymous (match any single dim, no binding).
     ``Scalar`` — scalar (no dimensions).
-    ``Value["expr"]`` — explicit runtime value expression for shape dims.
+    ``Value("expr")`` — explicit runtime value expression for shape dims.
 
 Tree structure symbols
     ``T``, ``S`` — named tree structure symbols.

src/shapix/_array_types.py

@@ -426,7 +426,7 @@ def _to_shape_spec(dims: tuple[object, ...]) -> tuple[DimSpec, ...]:
     else:
       msg = (
         "Invalid shape token "
-        f"{d!r}; expected int, Ellipsis, Dimension, Value[...], or a DimSpec"
+        f"{d!r}; expected int, Ellipsis, Dimension, Value(...), or a DimSpec"
       )
       raise TypeError(msg)
 

src/shapix/_dimensions.py

@@ -28,7 +28,7 @@ def flatten(x: F32[N, C]) -> F32[N * C]: ...
 Plain ``int`` values (e.g. ``3``) are also accepted as fixed dimension sizes
 when subscripting array types.
 
-Use ``Value["expr"]`` for dimensions that depend on runtime parameters or
+Use ``Value("expr")`` for dimensions that depend on runtime parameters or
 ``self`` attributes rather than previously bound shape names.
 """
 
@@ -190,11 +190,14 @@ def _dim_spec(self) -> DimSpec | None:
 
 
 class _ValueExpr:
-  """Runtime value expression used by ``Value[...]`` in shape subscripts."""
+  """Runtime value expression used by ``Value("...")`` in shape subscripts."""
 
   __slots__ = ("expr", "broadcastable")
 
-  def __init__(self, expr: str, *, broadcastable: bool = False) -> None:
+  def __init__(self, expr: object, *, broadcastable: bool = False) -> None:
+    if not isinstance(expr, str):
+      msg = "Value(...) expects a string expression"
+      raise TypeError(msg)
     self.expr = expr
     self.broadcastable = broadcastable
 
@@ -209,17 +212,20 @@ def _dim_spec(self) -> ValueDim:
 
   def __repr__(self) -> str:
     prefix = "+" if self.broadcastable else ""
-    return f'{prefix}Value["{self.expr}"]'
+    return f'{prefix}Value("{self.expr}")'
 
 
 # ---------------------------------------------------------------------------
 # Pre-defined dimension symbols
 # ---------------------------------------------------------------------------
 
 if tp.TYPE_CHECKING:
-  # Declared as ``Dimension`` so type checkers see the full operator set
-  # (``__add__``, ``__invert__``, ``__pos__``, …) and ``F32[N, C]`` subscripts.
+
   class Value:
+    def __new__(cls, expr: str) -> Dimension: ...
+
+    def __pos__(self) -> Dimension: ...
+
     @classmethod
     def __class_getitem__(cls, expr: str) -> Dimension: ...
 
@@ -236,20 +242,17 @@ def __class_getitem__(cls, expr: str) -> Dimension: ...
   __: Dimension
 else:
 
-  class Value:
+  class Value(_ValueExpr):
     """Explicit runtime value expression for shape annotations.
 
-    Use ``Value["size"]`` or ``Value["self.some_value + 3"]`` when a shape
+    Use ``Value("size")`` or ``Value("self.some_value + 3")`` when a shape
     depends on a runtime parameter rather than a previously bound dimension.
     """
 
     __slots__ = ()
 
-    def __class_getitem__(cls, expr: object) -> _ValueExpr:  # type: ignore[misc]
-      if not isinstance(expr, str):
-        msg = "Value[...] expects a string expression"
-        raise TypeError(msg)
-      return _ValueExpr(expr)
+    def __class_getitem__(cls, expr: object) -> Value:  # type: ignore[misc]
+      return cls(expr)
 
   # Common named dimensions
   Scalar = Dimension("")

src/shapix/_memo.py

@@ -187,7 +187,7 @@ def get_memo(_depth: int = 2) -> ShapeMemo:
 
 
 def get_scope(_depth: int = 2) -> dict[str, object]:
-  """Return the current runtime scope for ``Value[...]`` expressions."""
+  """Return the current runtime scope for ``Value(...)`` expressions."""
   explicit = _get_explicit_scope_stack()
   if explicit and explicit[-1] is not None:
     return explicit[-1]
@@ -201,7 +201,7 @@ def get_scope(_depth: int = 2) -> dict[str, object]:
 
   # beartype wrappers expose runtime arguments as ``args`` / ``kwargs`` plus
   # the wrapped function object. Rebind those to parameter names so
-  # ``Value["size"]`` and ``Value["self.attr"]`` work for both param and
+  # ``Value("size")`` and ``Value("self.attr")`` work for both param and
   # return validation.
   fn = locals_map.get("__beartype_func")
   args = locals_map.get("args")

src/shapix/_shape.py

@@ -13,7 +13,7 @@
 - :class:`SymbolicDim` — arithmetic expression evaluated against bound
   dimensions (``N+1``, ``H*W``). Optionally broadcastable.
 - :class:`ValueDim` — arithmetic expression evaluated against runtime call
-  scope (``Value["size"]``, ``Value["self.width + 3"]``). Optionally
+  scope (``Value("size")``, ``Value("self.width + 3")``). Optionally
   broadcastable.
 - :class:`VariadicDim` — matches zero or more contiguous dims and binds the
   matched sub-shape (``~spatial``). Optionally broadcastable.
@@ -96,7 +96,7 @@ class ValueDim:
 
   def __repr__(self) -> str:
     prefix = "+" if self.broadcastable else ""
-    return f'{prefix}Value["{self.expr}"]'
+    return f'{prefix}Value("{self.expr}")'
 
 
 @dataclass(frozen=True, slots=True)

src/shapix/_tree.py

@@ -105,7 +105,7 @@ def __call__(self, obj: object) -> bool:
     from ._memo import get_memo, get_scope, pop_memo, push_memo
 
     # Bridge memo + runtime scope so leaf checks reuse the caller's bindings
-    # and can resolve ``Value[...]`` expressions against the same parameters.
+    # and can resolve ``Value(...)`` expressions against the same parameters.
     memo = get_memo(_depth=3)
     scope = get_scope(_depth=3)
     push_memo(memo, scope=scope)

tests/test_dimensions.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import pytest
+
 from shapix._dimensions import Dimension, Value
 from shapix._shape import (
   ANONYMOUS,
@@ -222,14 +224,23 @@ def test_dim_plus_dim(self) -> None:
 
 
 class TestValueExpressions:
+  def test_value_expr_requires_string(self) -> None:
+    with pytest.raises(TypeError, match=r"Value\(\.\.\.\) expects a string expression"):
+      Value(1)  # type: ignore[arg-type]
+
   def test_value_expr(self) -> None:
-    spec = Value["size"]._dim_spec  # noqa: SLF001
+    spec = Value("size")._dim_spec  # noqa: SLF001
     assert isinstance(spec, ValueDim)
     assert spec.expr == "size"
     assert spec.broadcastable is False
 
   def test_broadcastable_value_expr(self) -> None:
-    spec = (+Value["size"])._dim_spec  # noqa: SLF001
+    spec = (+Value("size"))._dim_spec  # noqa: SLF001
     assert isinstance(spec, ValueDim)
     assert spec.expr == "size"
     assert spec.broadcastable is True
+
+  def test_value_getitem_alias(self) -> None:
+    spec = Value["size"]._dim_spec  # noqa: SLF001
+    assert isinstance(spec, ValueDim)
+    assert spec.expr == "size"

tests/test_numpy.py

@@ -418,8 +418,10 @@ def f(x: F32[N, C]) -> F32[N * C]:
 
 class TestValueExpressions:
   def test_value_expr_from_argument(self) -> None:
+    Size = Value("size")
+
     @beartype
-    def f(size: int, x: F32[Value["size"]]) -> F32[Value["size"]]:  # noqa: F821
+    def f(size: int, x: F32[Size]) -> F32[Size]:
       return x
 
     arr = np.ones(4, dtype=np.float32)
@@ -430,16 +432,19 @@ def f(size: int, x: F32[Value["size"]]) -> F32[Value["size"]]:  # noqa: F821
   def test_value_expr_from_self_attribute(self) -> None:
     class SomeClass:
       size = 5
+      FullSize = Value("self.size + 3")
 
       @beartype
-      def full(self) -> F32[Value["self.size + 3"]]:  # noqa: F821
+      def full(self) -> F32[FullSize]:
         return np.ones(self.size + 3, dtype=np.float32)
 
     assert SomeClass().full().shape == (8,)
 
   def test_value_expr_can_mix_scope_and_bound_dims(self) -> None:
+    Padded = Value("N + pad")
+
     @beartype
-    def f(x: F32[N], pad: int) -> F32[Value["N + pad"]]:  # noqa: F821
+    def f(x: F32[N], pad: int) -> F32[Padded]:
       return np.ones(x.shape[0] + pad, dtype=np.float32)
 
     assert f(np.ones(4, dtype=np.float32), 2).shape == (6,)

tests/typing/check_annotations.py

@@ -192,16 +192,20 @@ def double_channels(x: F32[N, C]) -> F32[N, C * 2]:
   return np.concatenate([x, x], axis=1).astype(np.float32)
 
 
+FromSize = Value("size")
+FromSelf = Value("self.offset + size")
+
+
 @beartype
-def from_value(size: int) -> F32[Value["size"]]:  # noqa: F821
+def from_value(size: int) -> F32[FromSize]:
   return np.ones(size, dtype=np.float32)
 
 
 class SomeClass:
   offset = 3
 
   @beartype
-  def from_self(self, size: int) -> F32[Value["self.offset + size"]]:  # noqa: F821
+  def from_self(self, size: int) -> F32[FromSelf]:
     return np.ones(self.offset + size, dtype=np.float32)