FIX: pickle nested unevaluated expressions (#509)

* DX: add regression tests for pickling nested expressions
* DX: ignore Ruff pickle rules in the test suite
* MAINT: use DataclassInstance in dataclass field helpers

Author: redeboer

pyproject.toml

@@ -545,6 +545,8 @@ split-on-trailing-comma = false
     "magic-value-comparison",
     "no-self-use",
     "private-member-access",
+    "suspicious-pickle-import",
+    "suspicious-pickle-usage",
     "too-many-locals",
     "too-many-positional-arguments",
     "unnecessary-collection-call",

src/ampform/sympy/_decorator.py

@@ -23,6 +23,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, Hashable, Iterable
 
+    from _typeshed import DataclassInstance
     from sympy.printing.latex import LatexPrinter
 
     if sys.version_info >= (3, 11):
@@ -279,7 +280,7 @@ def new_method(cls, *args, evaluate: bool = False, **kwargs) -> type[ExprClass]:
         return expr
 
     cls.__new__ = new_method
-    cls.__getnewargs__ = dataclasses.astuple
+    cls.__getnewargs__ = _get_field_values
     cls._hashable_content = _hashable_content_method
     if non_sympy_fields:
         cls._eval_subs = _eval_subs_method
@@ -486,7 +487,7 @@ def class_wrapper(cls: T) -> T:
 def _eval_subs_method(self, old, new, **hints):
     # https://github.com/sympy/sympy/blob/1.12/sympy/core/basic.py#L1117-L1147
     hit = False
-    old_args = dataclasses.astuple(self)
+    old_args = _get_field_values(self)
     new_args = list(old_args)
     for i, old_arg in enumerate(old_args):
         if not hasattr(old_arg, "_eval_subs"):
@@ -535,7 +536,7 @@ def _xreplace_method(self, rule) -> tuple[sp.Expr, bool]:
     if rule:
         new_args = []
         hit = False
-        for arg in dataclasses.astuple(self):
+        for arg in _get_field_values(self):
             if hasattr(arg, "_xreplace") and not isclass(arg):
                 replace_result, is_replaced = arg._xreplace(rule)  # ruff: ignore[private-member-access]
             elif isinstance(rule, abc.Mapping):
@@ -551,13 +552,28 @@ def _xreplace_method(self, rule) -> tuple[sp.Expr, bool]:
     return self, False
 
 
-def get_sympy_fields(cls) -> tuple[Field, ...]:
+def _get_field_values(self: DataclassInstance) -> tuple[Any, ...]:
+    """Get the field values of a dataclass-like class without recursing.
+
+    This is a shallow alternative to :func:`dataclasses.astuple`, which recurses into
+    nested dataclasses and converts them to `tuple` as well. Since decorated classes are
+    both dataclasses and `~sympy.core.expr.Expr` instances, that recursion would destroy
+    nested expressions, for instance when pickling.
+    """
+    return tuple(getattr(self, field.name) for field in dataclasses.fields(self))
+
+
+def get_sympy_fields(
+    cls: DataclassInstance | type[DataclassInstance],
+) -> tuple[Field[Any], ...]:
     return tuple(f for f in dataclasses.fields(cls) if _is_sympify(f))
 
 
-def get_non_sympy_fields(cls) -> tuple[Field, ...]:
+def get_non_sympy_fields(
+    cls: DataclassInstance | type[DataclassInstance],
+) -> tuple[Field[Any], ...]:
     return tuple(f for f in dataclasses.fields(cls) if not _is_sympify(f))
 
 
-def _is_sympify(field: Field) -> bool:
+def _is_sympify(field: Field[Any]) -> bool:
     return bool(field.metadata.get("sympify"))

tests/dynamics/test_deprecated.py

@@ -1,4 +1,4 @@
-import pickle  # ruff: ignore[suspicious-pickle-import]
+import pickle
 
 import sympy as sp
 
@@ -15,19 +15,19 @@ def test_pickle():
         # Pickle simple SymPy expression
         expr = z * angular_momentum
         pickled_obj = pickle.dumps(expr)
-        imported_expr = pickle.loads(pickled_obj)  # ruff: ignore[suspicious-pickle-usage]
+        imported_expr = pickle.loads(pickled_obj)
         assert expr == imported_expr
 
         # Pickle UnevaluatedExpression
         expr = UnevaluatedExpression()
         pickled_obj = pickle.dumps(expr)
-        imported_expr = pickle.loads(pickled_obj)  # ruff: ignore[suspicious-pickle-usage]
+        imported_expr = pickle.loads(pickled_obj)
         assert expr == imported_expr
 
         # Pickle classes derived from UnevaluatedExpression
         expr = BlattWeisskopfSquared(z, angular_momentum)
         pickled_obj = pickle.dumps(expr)
-        imported_expr = pickle.loads(pickled_obj)  # ruff: ignore[suspicious-pickle-usage]
+        imported_expr = pickle.loads(pickled_obj)
         assert expr == imported_expr
 
         expr = EnergyDependentWidth(
@@ -42,5 +42,5 @@ def test_pickle():
             name="Gamma_1",
         )
         pickled_obj = pickle.dumps(expr)
-        imported_expr = pickle.loads(pickled_obj)  # ruff: ignore[suspicious-pickle-usage]
+        imported_expr = pickle.loads(pickled_obj)
         assert expr == imported_expr

tests/helicity/test_helicity.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import pickle
 from typing import TYPE_CHECKING
 
 import pytest
@@ -170,6 +171,13 @@ def test_parameter_defaults_item_types(
             assert isinstance(symbol, sp.Symbol)
             assert isinstance(value, ParameterValue.__args__)
 
+    def test_pickle_roundtrip(self, amplitude_model: tuple[str, HelicityModel]):
+        """See https://github.com/ComPWA/ampform/issues/471."""
+        _, model = amplitude_model
+        pickled_model: HelicityModel = pickle.loads(pickle.dumps(model))
+        assert pickled_model.kinematic_variables == model.kinematic_variables
+        assert pickled_model == model
+
     def test_rename_symbols_no_renames(
         self, amplitude_model: tuple[str, HelicityModel]
     ):

tests/sympy/decorator/test_unevaluated.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import pickle
 from typing import Any, ClassVar
 
 import pytest
@@ -190,6 +191,32 @@ def evaluate(self) -> sp.Expr:
     assert isinstance(q_value.m2, sp.Float)
 
 
+@unevaluated(implement_doit=False)
+class _Outer(sp.Expr):
+    x: Any
+
+
+@unevaluated(implement_doit=False)
+class _Inner(sp.Expr):
+    x: Any
+    typ: type = argument(default=int, sympify=False)
+
+
+def test_pickle_nested_expressions():
+    x = sp.Symbol("x")
+    expr = _Outer(_Inner(x, typ=float))
+    pickled_expr: _Outer = pickle.loads(pickle.dumps(expr))
+    assert pickled_expr == expr
+    assert isinstance(pickled_expr.x, _Inner)
+    assert pickled_expr.x.x == x
+    assert pickled_expr.x.typ is float
+
+    sympifiable_expr = _Outer(_Outer(x))
+    pickled_expr = pickle.loads(pickle.dumps(sympifiable_expr))
+    assert pickled_expr == sympifiable_expr
+    assert isinstance(pickled_expr.x, _Outer)
+
+
 def test_subs_with_non_sympy_attributes():
     class Protocol: ...
 

tests/sympy/test_cache_helpers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-import pickle  # ruff: ignore[suspicious-pickle-import]
+import pickle
 import sys
 from concurrent.futures import ThreadPoolExecutor
 from threading import Event
@@ -64,7 +64,7 @@ def cached_function():
     assert not list(tmp_path.rglob("*.tmp"))
     cache_files = [path for path in tmp_path.rglob("*") if path.is_file()]
     assert len(cache_files) == 1
-    assert pickle.loads(cache_files[0].read_bytes()) == "result"  # ruff: ignore[suspicious-pickle-usage]
+    assert pickle.loads(cache_files[0].read_bytes()) == "result"
 
 
 @pytest.mark.parametrize("corrupt_data", [b"", b"not a pickle"])