From 0593694a9ef0083d62b758235b119410d7df053b Mon Sep 17 00:00:00 2001
From: Alistair Miles <alimanfoo@googlemail.com>
Date: Mon, 16 Sep 2024 16:25:49 +0100
Subject: [PATCH] Make ArrayWrapper pickleable (#425)

* make array wrapper pickleable

* fix requirements.txt

* whitespace
---
 allel/abc.py                            |  6 ++++++
 allel/test/model/ndarray/test_arrays.py | 24 +++++++++++++++++++++++-
 requirements.txt                        |  2 +-
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/allel/abc.py b/allel/abc.py
index c9192b31..f1865ad2 100644
--- a/allel/abc.py
+++ b/allel/abc.py
@@ -14,6 +14,12 @@ 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.
diff --git a/allel/test/model/ndarray/test_arrays.py b/allel/test/model/ndarray/test_arrays.py
index 44a3b0c4..05d79ab8 100644
--- a/allel/test/model/ndarray/test_arrays.py
+++ b/allel/test/model/ndarray/test_arrays.py
@@ -1,7 +1,11 @@
 # -*- coding: utf-8 -*-
+
+# stdlib imports
+import pickle
+import unittest
+
 # third-party imports
 import numpy as np
-import unittest
 import pytest
 
 
@@ -140,6 +144,12 @@ def test_no_subclass(self):
         self.assertNotIsInstance(g.T, GenotypeArray)
         self.assertNotIsInstance(g.astype('f4'), GenotypeArray)
 
+    def test_pickle(self):
+        g = self.setup_instance(diploid_genotype_data)
+        g2 = pickle.loads(pickle.dumps(g))
+        assert isinstance(g2, GenotypeArray)
+        aeq(g, g2)
+
 
 # noinspection PyMethodMayBeStatic
 class HaplotypeArrayTests(HaplotypeArrayInterface, unittest.TestCase):
@@ -208,6 +218,12 @@ def test_slice_types(self):
         assert isinstance(s, np.int8)
         assert not isinstance(s, HaplotypeArray)
 
+    def test_pickle(self):
+        h = self.setup_instance(haplotype_data)
+        h2 = pickle.loads(pickle.dumps(h))
+        assert isinstance(h2, HaplotypeArray)
+        aeq(h, h2)
+
 
 # noinspection PyMethodMayBeStatic
 class AlleleCountsArrayTests(AlleleCountsArrayInterface, unittest.TestCase):
@@ -283,6 +299,12 @@ def test_reduce_types(self):
             x = getattr(ac, m)(axis=1)
             assert isinstance(x, np.ndarray)
 
+    def test_pickle(self):
+        ac = self.setup_instance(allele_counts_data)
+        ac2 = pickle.loads(pickle.dumps(ac))
+        assert isinstance(ac2, AlleleCountsArray)
+        aeq(ac, ac2)
+
 
 # noinspection PyMethodMayBeStatic
 class GenotypeVectorTests(unittest.TestCase):
diff --git a/requirements.txt b/requirements.txt
index 106da8a1..6c7c0f73 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ scipy
 matplotlib
 seaborn
 pandas
-pomegranate
+protopunica
 scikit-learn
 h5py
 numexpr