pydiverse-common 0.3.4__py3-none-any.whl → 0.3.6__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"
 
@@ -103,7 +103,10 @@ class Dtype:
                 type_, pd.core.dtypes.common.classes(np_dtype)
             )
 
-        if pd.api.types.is_signed_integer_dtype(pandas_type):
+        workaround = (
+            pandas_type is not np.floating
+        )  # see https://github.com/pandas-dev/pandas/issues/62018
+        if workaround and pd.api.types.is_signed_integer_dtype(pandas_type):
             if is_np_dtype(pandas_type, np.int64):
                 return Int64()
             elif is_np_dtype(pandas_type, np.int32):
@@ -113,7 +116,7 @@ class Dtype:
             elif is_np_dtype(pandas_type, np.int8):
                 return Int8()
             raise TypeError
-        if pd.api.types.is_unsigned_integer_dtype(pandas_type):
+        if workaround and pd.api.types.is_unsigned_integer_dtype(pandas_type):
             if is_np_dtype(pandas_type, np.uint64):
                 return UInt64()
             elif is_np_dtype(pandas_type, np.uint32):
@@ -123,8 +126,8 @@ class Dtype:
             elif is_np_dtype(pandas_type, np.uint8):
                 return UInt8()
             raise TypeError
-        if pd.api.types.is_float_dtype(pandas_type):
-            if is_np_dtype(pandas_type, np.float64):
+        if not workaround or pd.api.types.is_float_dtype(pandas_type):
+            if not workaround or is_np_dtype(pandas_type, np.float64):
                 return Float64()
             elif is_np_dtype(pandas_type, np.float32):
                 return Float32()
@@ -141,6 +144,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 +210,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 +232,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 +300,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 +329,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 +476,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/util/structlog.py

@@ -80,7 +80,7 @@ class PydiverseConsoleRenderer(structlog.dev.ConsoleRenderer):
 
 def setup_logging(
     log_level=logging.INFO,
-    log_stream=sys.stderr,
+    log_stream=None,
     timestamp_format="%Y-%m-%d %H:%M:%S.%f",
 ):
     """Configures structlog and logging with sane defaults."""
@@ -92,6 +92,15 @@ def setup_logging(
             level=log_level,
             handlers=[StructlogHandler()],
         )
+        if log_stream is None:
+            try:
+                # hack to avoid dask pickling problems with pytest capture
+                import structlog._output
+
+                structlog._output.stderr = sys.stderr
+            finally:
+                pass
+            log_stream = sys.stderr
         # Configure structlog
         structlog.configure(
             processors=[

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydiverse-common
-Version: 0.3.4
+Version: 0.3.6
 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.6.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=LYZKaKYq_4uI4kUhoaCTTo5j1SRurswIOfN11Bkz25A,15986
 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
@@ -10,8 +10,8 @@ pydiverse/common/util/deep_merge.py,sha256=bV5p5_lsC-9nFah28EiEyG2h6U3Z5AuTqSoox
 pydiverse/common/util/disposable.py,sha256=4XoGz70YRWA9TAqnUBvRCTAdsOGBviFN0gzxU7veY9o,993
 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/util/structlog.py,sha256=xxhauxMuyxcKXTVg1MiPTkuvPBj8Zcr4o_v8Bq59Nig,3778
+pydiverse_common-0.3.6.dist-info/METADATA,sha256=TVaCnX9IArQwB64MGJ4wfVv1EE7KUDI51u9d0QnFXeY,3399
+pydiverse_common-0.3.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydiverse_common-0.3.6.dist-info/licenses/LICENSE,sha256=AcE6SDVuAq6v9ZLE_8eOCe_NvSE0rAPR3NR7lSowYh4,1517
+pydiverse_common-0.3.6.dist-info/RECORD,,