Replace ruff with black API for AST rendering (#80)

* introduce logger

* Fixed compatibility with pydiverse common 0.3.12

* prepare release 0.5.6

* replace ruff subprocess with black API

* Fix cast constraints. They should not depend on attributes.

Author: windiana42

docs/source/changelog.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 0.5.6 (2025-08-22)
+- improved regex behavior in .str.contains()
+- AST repr is printing user friendly syntax (python-like black formatted)
+- ColExpr repr is much more intuitive
+- fixed handling of const Enum (mutate(enum_col="value"))
+- added columns verb to simplify [c.name for c in tbl]
+- compatible with pydiverse.common 0.3.12
+
 ## 0.5.5 (2025-08-20)
 - don't suffix all joined columns if only the key columns overlap
 - add query to __str__/__repr__ operators of Table

pixi.lock

@@ -8,8 +8,9 @@ python = ">=3.10.18,<3.14"
 pandas = ">=2.3.0,<3"
 pyarrow = ">=20.0.0,<22"
 polars = ">=1.30.0,<2"
-pydiverse-common = ">=0.3.5,<0.4"
+pydiverse-common = ">=0.3.12,<0.4"
 sqlalchemy = ">=2.0.41,<3"
+black = ">=20"                     # used for printing ASTs
 
 [host-dependencies]
 pip = "*"
@@ -24,6 +25,7 @@ typos = "*"
 pixi-pycharm = ">=0.0.6"
 pytest = ">=7.1.2"
 pytest-xdist = ">=2.5.0"
+structlog = ">=25.4.0,<26"
 
 [feature.release.dependencies]
 hatch = ">=1.12.0"

pixi.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pydiverse-transform"
-version = "0.5.5"
+version = "0.5.6"
 description = "Pipe based dataframe manipulation library that can also transform data on SQL databases"
 authors = [
   { name = "QuantCo, Inc." },
@@ -30,8 +30,9 @@ dependencies = [
   "pandas >=2.3.0,<3",
   "pyarrow >=20.0.0,<22",
   "polars >=1.30.0,<2",
-  "pydiverse-common >=0.3.5,<0.4",
+  "pydiverse-common >=0.3.12,<0.4",
   "sqlalchemy >=2.0.41,<3",
+  "black >=20",
 ]
 
 [tool.hatch.build.targets.wheel]

pyproject.toml

@@ -608,6 +608,7 @@ def _is_in(x, *values):
 
     @impl(ops.str_contains)
     def _str_contains(x, y, allow_regex, true_if_regex_unsupported):
+        _ = true_if_regex_unsupported
         return x.str.contains(y, literal=not pl.select(allow_regex).item())
 
     @impl(ops.str_starts_with)

src/pydiverse/transform/_internal/backend/polars.py

@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import itertools
-import subprocess
 from collections.abc import Callable, Iterable
 from uuid import UUID
 
+from pydiverse.transform._internal.util.warnings import warn
+
 
 class AstNode:
     __slots__ = ["name"]
@@ -99,17 +100,13 @@ def next_nd(root: AstNode, cond: Callable[["AstNode"], bool]):
             + ")"
         )
         try:
-            proc = subprocess.run(
-                ["ruff", "format", "-"],  # '-' tells Ruff to read from stdin
-                input=unformatted.encode(),
-                capture_output=True,
-                check=True,
-            )
-        except subprocess.CalledProcessError as e:
-            print(unformatted)
-            print(e.stderr.decode())
-            raise e
-        return proc.stdout.decode()
+            import black
+
+            formatted = black.format_str(unformatted, mode=black.Mode(line_length=120))
+            return formatted
+        except Exception:
+            warn("Could not format AST representation with `black`.")
+            return unformatted
 
     def short_name(self) -> str:
         raise NotImplementedError()

src/pydiverse/transform/_internal/tree/ast.py

@@ -23,6 +23,7 @@
     Bool,
     Date,
     Datetime,
+    Decimal,
     Dtype,
     Duration,
     Enum,
@@ -2798,6 +2799,10 @@ def is_valid_cast(source, target) -> bool:
         }
 
         source = types.without_const(source)
+        if isinstance(source, String):
+            source = String()  # drop max_length
+        elif isinstance(source, Decimal):
+            source = Decimal()  # standardize precision/scale
 
         if isinstance(source, String) and isinstance(target, Enum):
             return True

src/pydiverse/transform/_internal/tree/col_expr.py

@@ -103,8 +103,28 @@ def converts_to(source: Dtype, target: Dtype) -> bool:
     source = without_const(source)
     if isinstance(source, List):
         return isinstance(target, List) and converts_to(source.inner, target.inner)
-    if isinstance(source, Enum):
-        return target == source or target == String()
+    if isinstance(source, Enum | String):
+        return (
+            target == source
+            or target == String()
+            or (
+                type(target) is String
+                and source.max_length is not None
+                and target.max_length > source.max_length
+            )
+        )
+    if isinstance(source, Decimal):
+        return (
+            target == source
+            or target in FLOAT_SUBTYPES
+            or target == Float()
+            or target == Decimal()
+            or (
+                isinstance(target, Decimal)
+                and target.scale >= source.scale
+                and (target.precision - target.scale >= source.precision - source.scale)
+            )
+        )
     return target in IMPLICIT_CONVS[source]
 
 
@@ -117,7 +137,7 @@ def to_python(dtype: Dtype):
         return float
     elif isinstance(dtype, List):
         return list
-    elif isinstance(dtype, Enum):
+    elif isinstance(dtype, Enum | String):
         return str
 
     return {
@@ -187,13 +207,23 @@ def lca_type(dtypes: list[Dtype]) -> Dtype:
 
         return List(lca_type([dtype.inner for dtype in dtypes]))
 
-    if any(isinstance(dtype, Enum) for dtype in dtypes):
+    if any(isinstance(dtype, Enum | String) for dtype in dtypes):
         if all(dtype == dtypes[0] for dtype in dtypes):
             return copy.copy(dtypes[0])
         if all(isinstance(dtype, Enum | String) for dtype in dtypes):
             return String()
         raise DataTypeError(f"incompatible types `{', '.join(str(d) for d in dtypes)}`")
 
+    if any(isinstance(dtype, Decimal) for dtype in dtypes):
+        if all(dtype == dtypes[0] for dtype in dtypes):
+            return copy.copy(dtypes[0])
+        if all(isinstance(dtype, Decimal) for dtype in dtypes):
+            precision_diff = max(dtype.precision - dtype.scale for dtype in dtypes)
+            scale = max(dtype.scale for dtype in dtypes)
+            precision = precision_diff + scale
+            return Decimal(precision, scale)
+        raise DataTypeError(f"incompatible types `{', '.join(str(d) for d in dtypes)}`")
+
     if not (
         common_ancestors := functools.reduce(
             operator.and_,
@@ -253,8 +283,12 @@ def is_subtype(dtype: Dtype) -> bool:
 def implicit_conversions(dtype: Dtype) -> list[Dtype]:
     if isinstance(dtype, List):
         return [List(inner) for inner in implicit_conversions(dtype.inner)]
-    if isinstance(dtype, Enum):
-        return [dtype, String()]
+    if isinstance(dtype, Enum | String):
+        return [String()] + ([dtype] if dtype.max_length is not None else [])
+    if isinstance(dtype, Decimal):
+        return (
+            list(FLOAT_SUBTYPES) + [Float()] + ([dtype] if dtype != Decimal() else [])
+        )
     return list(IMPLICIT_CONVS[dtype].keys())
 
 
@@ -303,8 +337,14 @@ def conversion_cost(dtype: Dtype, target: Dtype) -> tuple[int, int]:
     dtype = without_const(dtype)
     if isinstance(dtype, List):
         return conversion_cost(dtype.inner, target.inner)
-    if isinstance(dtype, Enum):
-        return (0, 0) if dtype == target else (0, 1)
+    if isinstance(dtype, Enum | String | Decimal):
+        return (
+            (0, 0)
+            if dtype == target
+            else (0, 1)
+            if type(dtype) is type(target)
+            else (0, 2)
+        )
     return IMPLICIT_CONVS[dtype][target]
 
 

src/pydiverse/transform/_internal/tree/types.py

@@ -1,8 +1,11 @@
 # Copyright (c) QuantCo and pydiverse contributors 2025-2025
 # SPDX-License-Identifier: BSD-3-Clause
+import logging
 
 import pytest
 
+from pydiverse.common.util.structlog import setup_logging
+
 # Setup
 
 
@@ -37,3 +40,6 @@ def pytest_collection_modifyitems(config: pytest.Config, items):
                     for kw in item.keywords
                 ):
                     item.add_marker(skip)
+
+
+setup_logging(log_level=logging.INFO)

tests/conftest.py

@@ -5,12 +5,15 @@
 
 import polars as pl
 import pytest
+import structlog
 
 import pydiverse.transform as pdt
 from pydiverse.transform._internal.errors import ColumnNotFoundError, DataTypeError
 from pydiverse.transform.extended import *
 from tests.util import assert_equal
 
+logger = structlog.get_logger(__name__)
+
 df1 = pl.DataFrame(
     {
         "col1": [1, 2, 3, 4],
@@ -845,6 +848,7 @@ def test_col_html_repr(self, tbl1):
 
     def test_expr_str(self, tbl1):
         expr_str = str(tbl1.col1 * 2)
+        logger.info(expr_str)
         assert "failed" not in expr_str
 
     def test_expr_html_repr(self, tbl1):
@@ -907,7 +911,7 @@ def test_verb_ast_repr(self, tbl3, tbl4):
             >> alias("tbl 42")
         )
 
-        (
+        ast_str = (
             intermed
             >> left_join(
                 tbl4
@@ -924,27 +928,23 @@ def test_verb_ast_repr(self, tbl3, tbl4):
 
         series = pl.Series([2**i for i in range(12)])
 
-        assert (
-            "tbl__1729.j"
-            in (
-                tbl3
-                >> mutate(
-                    u=C.col1 + 5,
-                    v=(tbl3.col2.exp() + tbl3.col4) * tbl3.col1,
-                    w=pdt.max(tbl3.col2, tbl3.col1, tbl3.col5.str.len()),
-                    x=pdt.count(),
-                    y=eval_aligned(
-                        (
-                            tbl3
-                            >> mutate(j=42)
-                            >> alias("tbl. 1729")
-                            >> filter(C.col2 > 0)
-                        ).j
-                        + tbl3.col1
-                    ),
-                    z=eval_aligned(series),
-                )
-                >> select(tbl3.col1, tbl3.col4)
-                >> left_join(tbl4, on="col1")
-            )._ast.ast_repr()
-        )
+        ast_str = (
+            tbl3
+            >> mutate(
+                u=C.col1 + 5,
+                v=(tbl3.col2.exp() + tbl3.col4) * tbl3.col1,
+                w=pdt.max(tbl3.col2, tbl3.col1, tbl3.col5.str.len()),
+                x=pdt.count(),
+                y=eval_aligned(
+                    (tbl3 >> mutate(j=42) >> alias("tbl. 1729") >> filter(C.col2 > 0)).j
+                    + tbl3.col1
+                ),
+                z=eval_aligned(series),
+            )
+            >> select(tbl3.col1, tbl3.col4)
+            >> left_join(tbl4, on="col1")
+        )._ast.ast_repr()
+
+        logger.info(f"AST:\n{ast_str}")
+
+        assert "tbl__1729.j" in ast_str