Repository: scikit-hep/awkward

Issue:
behaviors with decorated mixin methods cannot be pickled
### Version of Awkward Array

main

### Description and code to reproduce

The mixin machinery introduces a closure, which `pickle` can't serialise.


Discussion:

Resolving patch:
diff --git a/src/awkward/behaviors/mixins.py b/src/awkward/behaviors/mixins.py
--- a/src/awkward/behaviors/mixins.py
+++ b/src/awkward/behaviors/mixins.py
@@ -1,6 +1,7 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
 
 
+import functools
 import sys
 
 import awkward as ak
@@ -63,6 +64,10 @@ def register(cls):
     return register
 
 
+def _call_transposed(func, left, right):
+    return func(right, left)
+
+
 def mixin_class_method(ufunc, rhs=None, *, transpose=True):
     """
     Args:
@@ -70,7 +75,8 @@ def mixin_class_method(ufunc, rhs=None, *, transpose=True):
             hooked in Awkward Array, i.e. it can be the first argument of a behavior.
         rhs (Set[type] or None): Set of right-hand side argument types, optional
             if wrapping a unary function. The left-hand side is expected to
-            always be `self` of the parent class.
+            always be `self` of the parent class. The current class is implicitly
+            included in this set.
         transpose (bool): If true, automatically create a transpose signature
             (only makes sense for binary ufuncs).
 
@@ -86,12 +92,13 @@ def register(method):
                 ValueError("expected a set of right-hand-side argument types")
             )
         if transpose and rhs is not None:
-
-            def transposed(left, right):
-                return method(right, left)
-
             # make a copy of rhs, we will edit it later
-            method._awkward_mixin = (ufunc, set(rhs), transposed)
+            # use partial & a module-scoped function so that this is pickleable
+            method._awkward_mixin = (
+                ufunc,
+                set(rhs),
+                functools.partial(_call_transposed, method),
+            )
         else:
             method._awkward_mixin = (ufunc, rhs, None)
         return method


Test patch:
diff --git a/tests/test_2101_pickle_behavior_class.py b/tests/test_2101_pickle_behavior_class.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2101_pickle_behavior_class.py
@@ -0,0 +1,46 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import pickle
+
+import numpy as np
+import pytest  # noqa: F401
+
+import awkward as ak
+
+behavior = {}
+
+
+@ak.mixin_class(behavior, "vector")
+class Vector:
+    @ak.mixin_class_method(np.add, set())
+    def add(self, other):
+        return ak.zip(
+            {"x": self.x + other.x, "y": self.y + other.y, "z": self.z + other.z},
+            with_name="vector",
+        )
+
+
+def test():
+
+    vec = ak.Array(
+        [
+            {"x": 1, "y": 2, "z": 3},
+            {"x": 4, "y": 5, "z": 9},
+        ],
+        with_name="vector",
+        behavior=behavior,
+    )
+
+    assert ak._util.arrays_approx_equal(
+        (vec + vec),
+        ak.Array(
+            [
+                {"x": 2, "y": 4, "z": 6},
+                {"x": 8, "y": 10, "z": 18},
+            ],
+            with_name="vector",
+            behavior=behavior,
+        ),
+    )
+
+    assert ak._util.arrays_approx_equal(pickle.loads(pickle.dumps(vec)), vec)