From f9e17f71b395695d3d99c58f58c76a06c4ad8e12 Mon Sep 17 00:00:00 2001
From: Alistair Miles <alimanfoo@googlemail.com>
Date: Tue, 17 Sep 2024 11:22:15 +0100
Subject: [PATCH] properly fix unpickling recursion (#427)

---
 .gitignore                              |  3 +++
 allel/abc.py                            | 18 ++++++++----------
 allel/test/model/ndarray/test_arrays.py | 25 ++++++++++++++++++++++++-
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index 87ca1654..d19ffc83 100644
--- a/.gitignore
+++ b/.gitignore
@@ -69,3 +69,6 @@ temp/*
 
 # setuptools-scm
 allel/version.py
+
+# VScode
+.vscode
diff --git a/allel/abc.py b/allel/abc.py
index f1865ad2..56363e7b 100644
--- a/allel/abc.py
+++ b/allel/abc.py
@@ -14,12 +14,6 @@ def __init__(self, data):
             raise TypeError('values must be array-like')
         self._values = data
 
-    def __getstate__(self):
-        return self._values
-
-    def __setstate__(self, state):
-        self._values = state
-
     @property
     def values(self):
         """The underlying array of values.
@@ -40,10 +34,14 @@ def __repr__(self):
 
     def __getattr__(self, item):
         if item in {'__array_struct__', '__array_interface__'}:
-            # don't pass these through because we want to use __array__ to control numpy
-            # behaviour
-            raise AttributeError
-        return getattr(self.values, item)
+            # Don't pass these through because we want to use __array__ to control numpy
+            # behaviour.
+            raise AttributeError(item)
+        if item == "__setstate__":
+            # Special method called during unpickling, don't pass through.
+            raise AttributeError(item)
+        # Pass through all other attribute access to the wrapped values.
+        return getattr(self._values, item)
 
     def __getitem__(self, item):
         return self.values[item]
diff --git a/allel/test/model/ndarray/test_arrays.py b/allel/test/model/ndarray/test_arrays.py
index 05d79ab8..5d611d64 100644
--- a/allel/test/model/ndarray/test_arrays.py
+++ b/allel/test/model/ndarray/test_arrays.py
@@ -146,9 +146,18 @@ def test_no_subclass(self):
 
     def test_pickle(self):
         g = self.setup_instance(diploid_genotype_data)
-        g2 = pickle.loads(pickle.dumps(g))
+        s = pickle.dumps(g)
+        g2 = pickle.loads(s)
         assert isinstance(g2, GenotypeArray)
         aeq(g, g2)
+        gn = g2.to_n_alt()
+        assert isinstance(gn, np.ndarray)
+        ac = g2.count_alleles(max_allele=1)
+        assert isinstance(ac, AlleleCountsArray)
+        h = g2.to_haplotypes()
+        assert isinstance(h, HaplotypeArray)
+        gac = g2.to_allele_counts(max_allele=1)
+        assert isinstance(gac, GenotypeAlleleCountsArray)
 
 
 # noinspection PyMethodMayBeStatic
@@ -223,6 +232,10 @@ def test_pickle(self):
         h2 = pickle.loads(pickle.dumps(h))
         assert isinstance(h2, HaplotypeArray)
         aeq(h, h2)
+        ac = h2.count_alleles(max_allele=1)
+        assert isinstance(ac, AlleleCountsArray)
+        gt = h2.to_genotypes(ploidy=3)
+        assert isinstance(gt, GenotypeArray)
 
 
 # noinspection PyMethodMayBeStatic
@@ -392,6 +405,11 @@ def test_getitem(self):
         gs = gv[np.newaxis, :2, 0]  # change dimension semantics
         assert not isinstance(gs, GenotypeVector)
 
+    def test_pickle(self):
+        gv = GenotypeVector(diploid_genotype_data[0])
+        gv2 = pickle.loads(pickle.dumps(gv))
+        aeq(gv, gv2)
+
     # __getitem__
     # to_html_str
     # _repr_html_
@@ -497,3 +515,8 @@ def test_slice_types(self):
         s = g[0, 0, 0]
         assert isinstance(s, np.int8)
         assert not isinstance(s, GenotypeAlleleCountsArray)
+
+    def test_pickle(self):
+        g = GenotypeAlleleCountsArray(diploid_genotype_ac_data, dtype='i1')
+        g2 = pickle.loads(pickle.dumps(g))
+        aeq(g, g2)