Repository: hgrecco/pint-pandas

Issue:
PintArray fails to fit into a pickle jar
I am wondering if anyone has a quick workaround for pickling of `PintArray`. I am getting an error running following snippet (Python 3.8, pint-pandas 0.2):

```python
import pickle, pint_pandas
pickle.dumps(pint_pandas.PintArray([1.], dtype="m"))

# AttributeError: Can't pickle local object 'build_quantity_class.<locals>.Quantity'
```

It is also possible I am doing something totally wrong. Any ideas? Thx

Discussion:

Patch:
diff --git a/CHANGES b/CHANGES
index 746db0f..02d0c86 100644
--- a/CHANGES
+++ b/CHANGES
@@ -7,7 +7,7 @@ pint-pandas Changelog
 - Fixed a bug which caused the creation of PintArray (and even more so all binary mathematical operations)
   to convert input arrays to lists of Quantity, thereby causing a huge slowdown. #80
 - Fixed a bug where listlike operands in binary operations were not correctly converted to Quantity arrays. #86
-
+- Fixed pickling of `PintArray` instances. #71
 
 0.2 (2021-03-23)
 ----------------
diff --git a/pint_pandas/pint_array.py b/pint_pandas/pint_array.py
index 50505dd..9c418ca 100644
--- a/pint_pandas/pint_array.py
+++ b/pint_pandas/pint_array.py
@@ -205,6 +205,16 @@ class PintArray(ExtensionArray, ExtensionOpsMixin):
         self._data = values
         self._Q = self.dtype.ureg.Quantity
 
+    def __getstate__(self):
+        # we need to discard the cached _Q, which is not pickleable
+        ret = dict(self.__dict__)
+        ret.pop("_Q")
+        return ret
+
+    def __setstate__(self, dct):
+        self.__dict__.update(dct)
+        self._Q = self.dtype.ureg.Quantity
+
     @property
     def dtype(self):
         # type: () -> ExtensionDtype


Test patch:
diff --git a/pint_pandas/testsuite/test_issues.py b/pint_pandas/testsuite/test_issues.py
index 67b4ff4..787e24e 100644
--- a/pint_pandas/testsuite/test_issues.py
+++ b/pint_pandas/testsuite/test_issues.py
@@ -1,3 +1,4 @@
+import pickle
 import time
 
 import numpy as np
@@ -71,3 +72,10 @@ def test_issue_86():
     b_pa = PintArray([v.m_as(units) for v in b_listlike], units)
 
     assert np.all(a + b_listlike == a + b_pa)
+
+
+def test_issue_71():
+    a = PintArray([1, 2], ureg.m)
+    s = pickle.dumps(a)
+    b = pickle.loads(s)
+    assert np.all(a == b)