Fix unpickling of FieldArray instances in new interpreter contexts (#639)

* Allow pickling of galois fields

* Dynamically create and test pickled objects

* Add to changelog

* Make `FieldArray` subclasses hashable

* Improve unit tests

* Address @fyellin comments

---------

Co-authored-by: mhostetter <matthostetter@gmail.com>
Co-authored-by: mhostetter <mhostetter@users.noreply.github.com>

Author: fyellin

docs/release-notes/unreleased.md

@@ -12,11 +12,12 @@ tocdepth: 2
 
 ### Features
 
+- Make `FieldArray` subclasses hashable based on their properties. ([#639](https://github.com/mhostetter/galois/pull/639))
 - Improved FFT speed (~1.6x speedup) in `np.fft.fft()` with mixed-radix Cooley-Tukey algorithm. ([#620](https://github.com/mhostetter/galois/pull/620))
 
 ### Fixes
 
-- …
+- Fixed bug where `FieldArray` instances couldn't be unpickled if the `FieldArray` class had not yet be created. ([#639](https://github.com/mhostetter/galois/pull/639))
 
 ### Performance
 

src/galois/_fields/_factory.py

@@ -4,10 +4,11 @@
 
 from __future__ import annotations
 
+import copyreg
 import sys
 import types
 import warnings
-from typing import Type, overload
+from typing import Any, Dict, Tuple, Type, overload
 
 from typing_extensions import Literal
 
@@ -18,6 +19,7 @@
 from ..typing import PolyLike
 from ._array import FieldArray
 from ._gf2 import GF2
+from ._meta import FieldArrayMeta
 from ._primitive_element import is_primitive_element, primitive_element
 from ._ufunc import UFuncMixin_2_m, UFuncMixin_p_1, UFuncMixin_p_m
 
@@ -531,3 +533,46 @@ def _GF_extension(
 
 
 _GF_extension._classes = {}
+
+
+def _reconstruct_field_class(args: Tuple, kwargs: Dict[str, Any]):
+    """
+    Reconstruct a field class via `galois.GF(...)`.
+
+    Pickle's reduce protocol passes positional args only, so we wrap keyword arguments
+    in a dict and unpack them here.
+    """
+    return GF(*args, **kwargs)
+
+
+def _reduce_field_class(field_cls) -> Tuple[object, Tuple[Dict[str, Any]]]:
+    """
+    Pickle reducer for dynamically-created field classes (FieldArray subclasses).
+
+    We serialize the minimal set of constructor kwargs needed to reconstruct the same
+    field class via the GF factory on unpickle.
+    """
+    args = (
+        int(field_cls.characteristic),
+        int(field_cls.degree),
+    )
+
+    kwargs: Dict[str, Any] = {
+        "primitive_element": int(field_cls.primitive_element),
+        "verify": False,
+        "compile": field_cls.ufunc_mode,  # Restore the field's current ufunc mode on reconstruction
+        "repr": field_cls.element_repr,
+    }
+
+    # Only extension fields have an irreducible polynomial. Encode as a string to avoid
+    # formatting / parsing issues.
+    if field_cls.degree > 1:
+        kwargs["irreducible_poly"] = str(field_cls.irreducible_poly)
+
+    # Return (callable, args) where args is a tuple of positional args; we pass kwargs as one arg.
+    return (_reconstruct_field_class, (args, kwargs))
+
+
+# Register pickling for the metaclass used by field classes.
+# FieldArrayMeta is your metaclass (import it appropriately here).
+copyreg.pickle(FieldArrayMeta, _reduce_field_class)

src/galois/_fields/_meta.py

@@ -51,6 +51,10 @@ def __init__(cls, name, bases, namespace, **kwargs):
         cls._name = f"GF({cls._order_str})"
         cls._long_name = f"GF({cls._order_str}, primitive_element={cls._primitive_element_str!r}, irreducible_poly={cls._irreducible_poly_str!r})"
 
+    def __hash__(cls):
+        t = (cls.characteristic, cls.degree, cls._irreducible_poly_int, int(cls._primitive_element))
+        return hash(t)
+
     def __repr__(cls) -> str:
         if cls.order == 0:
             # This is not a runtime-created subclass, so return the base class name.

src/galois/_polys/_poly.py

@@ -1014,7 +1014,7 @@ def __int__(self) -> int:
         return self.__index__()
 
     def __hash__(self):
-        t = (self.field.order, *self.nonzero_degrees.tolist(), *self.nonzero_coeffs.tolist())
+        t = (hash(self.field), hash(tuple(self.nonzero_degrees.tolist())), hash(tuple(self.nonzero_coeffs.tolist())))
         return hash(t)
 
     @overload

tests/fields/test_classes.py

@@ -2,8 +2,6 @@
 A pytest module to test the class attributes of FieldArray subclasses.
 """
 
-import pickle
-
 import numpy as np
 import pytest
 
@@ -117,37 +115,3 @@ def test_is_primitive_poly():
     poly = galois.conway_poly(3, 101)
     GF = galois.GF(3**101, irreducible_poly=poly, primitive_element="x", verify=False)
     assert GF.is_primitive_poly
-
-
-def test_pickle_class(tmp_path):
-    GF = galois.GF(13)
-    with open(tmp_path / "class.pkl", "wb") as f:
-        pickle.dump(GF, f)
-    with open(tmp_path / "class.pkl", "rb") as f:
-        GF_loaded = pickle.load(f)
-    assert GF is GF_loaded
-
-    GF = galois.GF(3**5)
-    with open(tmp_path / "class.pkl", "wb") as f:
-        pickle.dump(GF, f)
-    with open(tmp_path / "class.pkl", "rb") as f:
-        GF_loaded = pickle.load(f)
-    assert GF is GF_loaded
-
-
-def test_pickle_array(tmp_path):
-    GF = galois.GF(13)
-    x = GF.Random(10)
-    with open(tmp_path / "array.pkl", "wb") as f:
-        pickle.dump(x, f)
-    with open(tmp_path / "array.pkl", "rb") as f:
-        x_loaded = pickle.load(f)
-    assert np.array_equal(x, x_loaded)
-
-    GF = galois.GF(3**5)
-    x = GF.Random(10)
-    with open(tmp_path / "array.pkl", "wb") as f:
-        pickle.dump(x, f)
-    with open(tmp_path / "array.pkl", "rb") as f:
-        x_loaded = pickle.load(f)
-    assert np.array_equal(x, x_loaded)

tests/fields/test_pickle.py

@@ -0,0 +1,80 @@
+"""
+A pytest module to test pickling and unpickling of FieldArray subclasses and their instances.
+"""
+
+import os
+import pickle
+import subprocess
+import sys
+
+import numpy as np
+import pytest
+
+import galois
+
+
+@pytest.mark.parametrize("order", [3, 3**2])
+def test_pickle_field_array_in_same_interpreter(order, tmp_path):
+    GF = galois.GF(order)
+    x = GF.Random(10)
+
+    # Write the pickle artifact
+    pkl_path = tmp_path / "field_array.pkl"
+    with pkl_path.open("wb") as f:
+        pickle.dump(x, f)
+
+    # Read the pickle artifact
+    with pkl_path.open("rb") as f:
+        x2 = pickle.load(f)
+
+    assert type(x) is type(x2)
+    assert np.array_equal(x, x2)
+
+
+@pytest.mark.parametrize("order", [3, 3**2])
+def test_pickle_field_array_in_new_interpreter(order, tmp_path):
+    GF = galois.GF(order)
+    x = GF.Random(10)
+
+    # Write the pickle artifact
+    pkl_path = tmp_path / "field_array.pkl"
+    with pkl_path.open("wb") as f:
+        pickle.dump(x, f)
+
+    # Capture "expected" in a representation that is easy to embed as Python literals.
+    # Use integers of the underlying representation, not repr(x).
+    expected_hash = hash(type(x))
+    expected_properties = type(x).properties
+    expected_values = x.tolist()
+
+    # Run a fresh interpreter that ONLY imports galois and unpickles.
+    code = f"""
+import pickle
+import numpy as np
+import galois
+
+with open(r"{pkl_path}", "rb") as f:
+    x2 = pickle.load(f)
+
+assert hash(type(x2)) == {expected_hash}
+assert type(x2).properties == {expected_properties!r}
+assert x2.tolist() == {expected_values!r}
+"""
+
+    env = os.environ.copy()
+
+    # Ensure subprocess imports your working tree version (editable checkouts, etc.)
+    repo_root = os.getcwd()
+    env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
+
+    result = subprocess.run(
+        [sys.executable, "-c", code],
+        env=env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+    assert result.returncode == 0, (
+        f"Unpickling in a fresh interpreter failed.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}\n"
+    )