pydiverse-common 0.3.4__py3-none-any.whl → 0.3.5__py3-none-any.whl

pydiverse/common/__init__.py

@@ -7,6 +7,7 @@ from .dtypes import (
     Decimal,
     Dtype,
     Duration,
+    Enum,
     Float,
     Float32,
     Float64,
@@ -51,5 +52,6 @@ __all__ = [
     "UInt32",
     "UInt64",
     "List",
+    "Enum",
     "PandasBackend",
 ]

pydiverse/common/dtypes.py

@@ -1,9 +1,9 @@
 # Copyright (c) QuantCo and pydiverse contributors 2025-2025
 # SPDX-License-Identifier: BSD-3-Clause
-from enum import Enum
+import enum
 
 
-class PandasBackend(str, Enum):
+class PandasBackend(str, enum.Enum):
     NUMPY = "numpy"
     ARROW = "arrow"
 
@@ -141,6 +141,9 @@ class Dtype:
         # we don't know any decimal/time/null dtypes in pandas if column is not
         # arrow backed
 
+        if pandas_type.name == "category":
+            return Enum(*pandas_type.categories.to_list())
+
         raise TypeError
 
     @staticmethod
@@ -204,6 +207,8 @@ class Dtype:
 
         if isinstance(polars_type, pl.List):
             return List(Dtype.from_polars(polars_type.inner))
+        if isinstance(polars_type, pl.Enum):
+            return Enum(*polars_type.categories)
 
         return {
             pl.Int64: Int64(),
@@ -224,7 +229,6 @@ class Dtype:
             pl.Date: Date(),
             pl.Null: NullType(),
             pl.Duration: Duration(),
-            pl.Enum: String(),
         }[polars_type.base_type()]
 
     def to_sql(self):
@@ -293,6 +297,9 @@ class Dtype:
         if isinstance(self, List):
             raise TypeError("pandas doesn't have a native list dtype")
 
+        if isinstance(self, Enum):
+            return pd.CategoricalDtype(self.categories)
+
         return {
             Int(): pd.Int64Dtype(),  # we default to 64 bit
             Int8(): pd.Int8Dtype(),
@@ -319,6 +326,9 @@ class Dtype:
         """Convert this Dtype to a PyArrow type."""
         import pyarrow as pa
 
+        if isinstance(self, Enum):
+            return pa.string()
+
         return {
             Int(): pa.int64(),  # we default to 64 bit
             Int8(): pa.int8(),
@@ -463,3 +473,36 @@ class List(Dtype):
         import pyarrow as pa
 
         return pa.list_(self.inner.to_arrow())
+
+
+class Enum(String):
+    def __init__(self, *categories: str):
+        if not all(isinstance(c, str) for c in categories):
+            raise TypeError("arguments for `Enum` must have type `str`")
+        self.categories = list(categories)
+
+    def __eq__(self, rhs):
+        return isinstance(rhs, Enum) and self.categories == rhs.categories
+
+    def __repr__(self) -> str:
+        return f"Enum[{', '.join(repr(c) for c in self.categories)}]"
+
+    def __hash__(self):
+        return hash(tuple(self.categories))
+
+    def to_polars(self):
+        import polars as pl
+
+        return pl.Enum(self.categories)
+
+    def to_sql(self):
+        import sqlalchemy as sqa
+
+        return sqa.String()
+
+    def to_arrow(self):
+        import pyarrow as pa
+
+        # There is also pa.dictionary(), which seems to be kind of similar to an enum.
+        # Maybe it is better to convert to this.
+        return pa.string()

{pydiverse_common-0.3.4.dist-info → pydiverse_common-0.3.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydiverse-common
-Version: 0.3.4
+Version: 0.3.5
 Summary: Common functionality shared between pydiverse libraries
 Author: QuantCo, Inc.
 Author-email: Martin Trautmann <windiana@users.sf.net>, Finn Rudolph <finn.rudolph@t-online.de>
@@ -48,7 +48,7 @@ Classifier: Topic :: Database
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
 Requires-Python: <3.14,>=3.10
-Requires-Dist: python-box<8.0.0,>=7.3.2
+Requires-Dist: python-box<8,>=7.3.2
 Description-Content-Type: text/markdown
 
 # pydiverse.common

{pydiverse_common-0.3.4.dist-info → pydiverse_common-0.3.5.dist-info}/RECORD

@@ -1,5 +1,5 @@
-pydiverse/common/__init__.py,sha256=eK3Wji6868cvXeq6e9YPRxi20vU6i9Lql4IzG6iB4B4,784
-pydiverse/common/dtypes.py,sha256=1ADpaWJbdNIFEE2DFjdIy2RRsp5Wg3UVnSM1ZN8XVdw,14555
+pydiverse/common/__init__.py,sha256=J7b4iStFyaEMYre_jdlZ4l_8dLyrMWCIpQdsMQcB8aI,806
+pydiverse/common/dtypes.py,sha256=Twna-Esfy0-4wzIxX37C1CzPTJz8MnesV_LNegndKzE,15787
 pydiverse/common/testing.py,sha256=FcivI5wn0X3gzJhwnysKvCOgjSTTXaN6FtSFJ72jfSg,341
 pydiverse/common/version.py,sha256=1IU_m4r76_Qq0u-Tyo2_bERZFOkh0ZFueVzDqcCfLO0,336
 pydiverse/common/errors/__init__.py,sha256=FNeEfVbUa23b9sHkFsmxHYhY6sRgjaZysPQmlovpJrI,262
@@ -11,7 +11,7 @@ pydiverse/common/util/disposable.py,sha256=4XoGz70YRWA9TAqnUBvRCTAdsOGBviFN0gzxU
 pydiverse/common/util/hashing.py,sha256=6x77BKg-w61u59fuTe9di0BtU-kEKH6UTRcKsRoYJ84,1196
 pydiverse/common/util/import_.py,sha256=K7dSgz4YyrqEvqhoOzbwgD7D8HScMoO5XoSWtjbaoUs,4056
 pydiverse/common/util/structlog.py,sha256=g0d8yaXBzAxmGNGZYMnMP9dsSQ__jN44GAY8Mb0ABeI,3487
-pydiverse_common-0.3.4.dist-info/METADATA,sha256=iV-aHd6IDEbgTaIsbB5lgAstdXaAAdotwZCh05laHA4,3403
-pydiverse_common-0.3.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydiverse_common-0.3.4.dist-info/licenses/LICENSE,sha256=AcE6SDVuAq6v9ZLE_8eOCe_NvSE0rAPR3NR7lSowYh4,1517
-pydiverse_common-0.3.4.dist-info/RECORD,,
+pydiverse_common-0.3.5.dist-info/METADATA,sha256=E4Z7zT2kwRtIlGetzPlxFf5Oj5-b037lAkLikUbNgKk,3399
+pydiverse_common-0.3.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydiverse_common-0.3.5.dist-info/licenses/LICENSE,sha256=AcE6SDVuAq6v9ZLE_8eOCe_NvSE0rAPR3NR7lSowYh4,1517
+pydiverse_common-0.3.5.dist-info/RECORD,,