Monkeypatch PandasArray class (#22)

childmindresearch · Feb 14, 2024 · a15a1aa · a15a1aa
1 parent c8cb4bc
commit a15a1aa
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 3 deletions.
diff --git a/elbow/dtypes/_json.py b/elbow/dtypes/_json.py
@@ -4,8 +4,8 @@
 import numpy as np
 import pyarrow as pa
 from pandas.api.extensions import register_extension_dtype
-from pandas.core.arrays import PandasArray
 
+from ._pandas_array import PandasArray
 from .base import PaExtensionArray, PaExtensionScalar, PaExtensionType, PdExtensionDtype
 
 __all__ = [

diff --git a/elbow/dtypes/_ndarray.py b/elbow/dtypes/_ndarray.py
@@ -3,8 +3,8 @@
 import numpy as np
 import pyarrow as pa
 from pandas.api.extensions import register_extension_dtype
-from pandas.core.arrays import PandasArray
 
+from ._pandas_array import PandasArray
 from .base import PaExtensionArray, PaExtensionScalar, PaExtensionType, PdExtensionDtype
 
 __all__ = [

diff --git a/elbow/dtypes/_pandas_array.py b/elbow/dtypes/_pandas_array.py
@@ -0,0 +1,13 @@
+"""Monkeypatch PandasArray class.
+
+Pandas 2.x renamed and moved the PandasArray class.
+"""
+
+try:
+    # Pandas >= 2.x
+    from pandas.arrays import NumpyExtensionArray as PandasArray
+except ImportError:
+    # Pandas < 2.0
+    from pandas.core.arrays import PandasArray
+
+__all__ = ["PandasArray"]
diff --git a/elbow/dtypes/_pickle.py b/elbow/dtypes/_pickle.py
@@ -4,8 +4,8 @@
 import numpy as np
 import pyarrow as pa
 from pandas.api.extensions import register_extension_dtype
-from pandas.core.arrays import PandasArray
 
+from ._pandas_array import PandasArray
 from .base import PaExtensionArray, PaExtensionScalar, PaExtensionType, PdExtensionDtype
 
 __all__ = [