Fix: Error when using ak.copy in v2 (#1532)

* Fix: Error when using ak.copy in v2 * Adding tests for Record deepcopy() * Accept a `memo` argument, as `__deepcopy__` expects. * The keys of `memo` should be `id(old_instance)` and the values should be `new_instance`. * Future-proof from smarter linters that would probably complain about `setattr` when the key is known. Co-authored-by: Jim Pivarski <[email protected]>
scikit-hep · Jul 8, 2022 · 818b373 · 818b373
1 parent 2885a0e
commit 818b373
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 9 deletions.
diff --git a/src/awkward/_v2/contents/content.py b/src/awkward/_v2/contents/content.py
@@ -1480,15 +1480,13 @@ def with_parameter(self, key, value):
 
         return out
 
-    def deep_copy(self, memo=None):
+    def __deepcopy__(self, memo=None):
         cls = self.__class__
         new_instance = cls.__new__(cls)
-        memo = {"id(self)": new_instance}
+        memo = {id(self): new_instance}
         for k, v in self.__dict__.items():
             if k == "_nplike":
-                setattr(new_instance, k, v)
-            # elif isinstance(v, ak._v2.index.Index):
-            #     setattr(new_instance, k, copy.copy(v))
+                new_instance._nplike = v
             else:
                 setattr(new_instance, k, copy.deepcopy(v, memo))
         return new_instance

diff --git a/src/awkward/_v2/operations/ak_copy.py b/src/awkward/_v2/operations/ak_copy.py
@@ -63,4 +63,4 @@ def _impl(array):
         allow_record=True,
         allow_other=False,
     )
-    return ak._v2._util.wrap(layout.deep_copy(), ak._v2._util.behavior_of(array))
+    return ak._v2._util.wrap(layout.__deepcopy__(), ak._v2._util.behavior_of(array))
diff --git a/src/awkward/_v2/record.py b/src/awkward/_v2/record.py
@@ -1,6 +1,5 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
 
-import copy
 from collections.abc import Iterable
 
 import awkward as ak
@@ -191,8 +190,8 @@ def to_list(self, behavior=None):
 
         return self._array[self._at : self._at + 1]._to_list(behavior, None)[0]
 
-    def deep_copy(self):
-        return Record(self._array.deep_copy(), copy.deepcopy(self._at))
+    def __deepcopy__(self, memo=None):
+        return Record(self._array.__deepcopy__(memo), self._at)
 
     def recursively_apply(
         self,

diff --git a/tests/v2/test_0511-copy-and-deepcopy.py b/tests/v2/test_0511-copy-and-deepcopy.py
@@ -53,3 +53,39 @@ def test():
         {"x": 3, "y": 9},
     ]
     assert to_list(original) == [{"x": 1}, {"x": 2}, {"x": 3}]
+
+    array = ak._v2.Array(
+        [
+            [{"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}, {"x": 3.3, "y": [1, 2, 3]}],
+            [],
+            [{"x": 4.4, "y": [1, 2, 3, 4]}, {"x": 5.5, "y": [1, 2, 3, 4, 5]}],
+        ]
+    )
+
+    assert to_list(ak._v2.operations.copy(array)) == [
+        [{"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}, {"x": 3.3, "y": [1, 2, 3]}],
+        [],
+        [{"x": 4.4, "y": [1, 2, 3, 4]}, {"x": 5.5, "y": [1, 2, 3, 4, 5]}],
+    ]
+
+    a = ak._v2.Array(
+        [
+            {"x": 0, "y": 0.0},
+            {"x": 1, "y": 1.1},
+            {"x": 2, "y": 2.2},
+            {"x": 3, "y": 3.3},
+            {"x": 4, "y": 4.4},
+        ]
+    )
+    record = a[2].layout.__deepcopy__()
+    a["z"] = a.x**2
+
+    assert to_list(a) == [
+        {"x": 0, "y": 0.0, "z": 0},
+        {"x": 1, "y": 1.1, "z": 1},
+        {"x": 2, "y": 2.2, "z": 4},
+        {"x": 3, "y": 3.3, "z": 9},
+        {"x": 4, "y": 4.4, "z": 16},
+    ]
+    assert to_list(record) == {"x": 2, "y": 2.2}
+    assert to_list(a[2]) == {"x": 2, "y": 2.2, "z": 4}