Fixes for mypy + Enumeration

equinox/_enum.py

@@ -59,9 +59,15 @@ def __new__(mcs, cls_name, bases, namespace):
             _base_offsets=base_offsets,
         )
         if cls_name == "Enumeration" and "Enumeration" not in globals().keys():
+            new_namespace["string"] = namespace.pop("string")
             new_namespace["promote"] = namespace.pop("promote")
             new_namespace["where"] = namespace.pop("where")
         else:
+            if "string" in namespace:
+                raise ValueError(
+                    "Cannot have enumeration item with name `string`, as this "
+                    "conflicts with the classmethod of this name"
+                )
             if "promote" in namespace:
                 raise ValueError(
                     "Cannot have enumeration item with name `promote`, as this "
@@ -207,16 +213,17 @@ def is_traced(self) -> bool:
     from typing import ClassVar
     from typing_extensions import Self
 
-    class _Sequence(type):
-        def __getitem__(cls, item) -> str: ...
-
+    class _HasLen(type):
         def __len__(cls) -> int: ...
 
-    class Enumeration(enum.Enum, EnumerationItem, metaclass=_Sequence):  # pyright: ignore
+    class Enumeration(enum.Enum, EnumerationItem, metaclass=_HasLen):  # pyright: ignore
         _name_to_item: ClassVar[dict[str, EnumerationItem]]  # pyright: ignore
         _index_to_message: ClassVar[list[str]]  # pyright: ignore
         _base_offsets: ClassVar[dict["Enumeration", int]]
 
+        @classmethod
+        def string(cls, item: "Enumeration") -> str: ...
+
         @classmethod
         def promote(cls, item: "Enumeration") -> Self: ...
 
@@ -262,10 +269,10 @@ class RESULTS(eqx.Enumeration):
         print(RESULTS.success)  # RESULTS<Hurrah!>
         ```
 
-        Given a Enumeration element, just the string can be looked up by indexing it:
+        Given an Enumeration element, just the string message can be looked up:
         ```python
         result = RESULTS.success
-        print(RESULTS[result])  # Hurrah!
+        print(RESULTS.string(result))  # Hurrah!
         ```
 
         Enumerations support inheritance, to include all of the superclasses' fields as
@@ -292,6 +299,37 @@ class MORE_RESULTS(RESULTS):  # pyright: ignore
         string corresponding to the enumeration item.
         """
 
+        @classmethod
+        def string(cls, item: _Enumeration) -> str:
+            """Look up the string associated with an enum element.
+
+            !!! Example
+
+                ```python
+                class RESULTS(eqx.Enumeration):
+                    success = "success"
+                    linear_solve_failed = "Linear solve failed to converge"
+
+                result = RESULTS.linear_solve_failed
+                assert RESULTS.string(result) == "Linear solve failed to converge"
+                ```
+
+            **Arguments:**
+
+            - `item`: a member of the enumeration.
+
+            **Returns:**
+
+            If `item` is outside of JIT, then the string associated with the enum (from
+            the right hand side of the equals sign during its definition).
+
+            If `item` is a traced JAX value (inside of JIT), then the string "traced" is
+            returned.
+            """
+            # Wrap old API with new API, see
+            # https://github.com/patrick-kidger/equinox/issues/1121.
+            return cls[item]
+
         @classmethod
         def promote(cls, item: _Enumeration) -> _Enumeration:
             """Enums support `.promote` (on the class) to promote from an inherited

equinox/_filters.py

@@ -22,7 +22,9 @@
 _array_types = (np.ndarray, np.generic, jax.Array)  # JAX < 0.7.2
 
 try:  # JAX 0.7.2
-    from jax._src.literals import LiteralArray
+    from jax._src.literals import (
+        LiteralArray,  # pyright: ignore[reportAttributeAccessIssue]
+    )
 
     _array_types += (LiteralArray,)
 except ImportError:

equinox/_serialisation.py

@@ -135,7 +135,7 @@ def default_deserialise_filter_spec(f: BinaryIO, x: Any) -> Any:
         out = np.load(f)
         if isinstance(x, jax.dtypes.bfloat16):
             out = out.view(jax.dtypes.bfloat16)
-        return type(x)(out.item())
+        return type(x)(out.item())  # pyright: ignore[reportCallIssue]
     else:
         return x
 

tests/test_enum.py

@@ -207,13 +207,27 @@ class A(eqxi.Enumeration):
     class B(eqxi.Enumeration):
         b = "bye"
 
-    A[A.a]
+    assert A[A.a] == "hi"
     with pytest.raises(ValueError):
         A[0]
     with pytest.raises(ValueError):
         A[B.b]
 
 
+def test_string():
+    class A(eqxi.Enumeration):
+        a = "hi"
+
+    class B(eqxi.Enumeration):
+        b = "bye"
+
+    assert A.string(A.a) == "hi"
+    with pytest.raises(ValueError):
+        A.string(0)  # pyright: ignore[reportArgumentType]
+    with pytest.raises(ValueError):
+        A.string(B.b)
+
+
 def test_isinstance():
     class A(eqxi.Enumeration):
         a = "hi"

tests/test_serialisation.py

@@ -122,22 +122,22 @@ def test_helpful_errors(getkey, tmp_path):
 
 def test_generic_dtype_serialisation(getkey, tmp_path):
     # Ensure we can round trip when we start with an array
-    jax_array = jnp.array(bfloat16(1))
+    jax_array = jnp.array(bfloat16(1))  # pyright: ignore[reportCallIssue]
     eqx.tree_serialise_leaves(tmp_path, jax_array)
-    like_jax_array = jnp.array(bfloat16(2))
+    like_jax_array = jnp.array(bfloat16(2))  # pyright: ignore[reportCallIssue]
     loaded_jax_array = eqx.tree_deserialise_leaves(tmp_path, like_jax_array)
     assert jax_array.item() == loaded_jax_array.item()
 
     tree = (
         jnp.array(1e-8),
-        bfloat16(1e-8),
+        bfloat16(1e-8),  # pyright: ignore[reportCallIssue]
         np.float32(1e-8),
         jnp.array(1e-8),
         np.float64(1e-8),
     )
     like_tree = (
         jnp.array(2.0),
-        bfloat16(2),
+        bfloat16(2),  # pyright: ignore[reportCallIssue]
         np.float32(2),
         jnp.array(2.0),
         np.float64(2.0),

tests/test_tree.py

@@ -238,9 +238,9 @@ def test_tree_equal_scalars():
     # dtype does matter
     assert _typeequal(eqx.tree_equal(x, z), False)
 
-    z = jax.dtypes.bfloat16(1)
-    z2 = jax.dtypes.bfloat16(1)
-    w = jax.dtypes.bfloat16(2)
+    z = jax.dtypes.bfloat16(1)  # pyright: ignore[reportCallIssue]
+    z2 = jax.dtypes.bfloat16(1)  # pyright: ignore[reportCallIssue]
+    w = jax.dtypes.bfloat16(2)  # pyright: ignore[reportCallIssue]
     assert _typeequal(eqx.tree_equal(z, z2), True)
     assert _typeequal(eqx.tree_equal(z, w), False)