Skip to content

Commit 83ebb72

Browse files
timsaucerclaude
andcommitted
refactor: add coerce_to_expr helpers and replace inline coercion patterns
Introduce coerce_to_expr() and coerce_to_expr_or_none() in expr.py as the complement to ensure_expr() — where ensure_expr rejects non-Expr values, these helpers wrap them via Expr.literal(). Replaces ~60 inline isinstance checks in functions.py with single-line helper calls, and updates the make-pythonic skill to document the new pattern. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 7515ebc commit 83ebb72

3 files changed

Lines changed: 167 additions & 169 deletions

File tree

.ai/skills/make-pythonic/SKILL.md

Lines changed: 33 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -261,28 +261,46 @@ python -m pytest --doctest-modules python/datafusion/functions.py -v
261261

262262
## Coercion Helper Pattern
263263

264-
When adding coercion to a function, use this inline pattern:
264+
Use the coercion helpers from `datafusion.expr` to convert native Python values to `Expr`. These are the complement of `ensure_expr()` — where `ensure_expr` *rejects* non-`Expr` values, the coercion helpers *wrap* them via `Expr.literal()`.
265+
266+
**For required parameters** use `coerce_to_expr`:
265267

266268
```python
267-
if not isinstance(arg, Expr):
268-
arg = Expr.literal(arg)
269+
from datafusion.expr import coerce_to_expr
270+
271+
def left(string: Expr, n: Expr | int) -> Expr:
272+
n = coerce_to_expr(n)
273+
return Expr(f.left(string.expr, n.expr))
269274
```
270275

271-
Do NOT create a new helper function for this — the inline check is clear and explicit. The project intentionally uses `ensure_expr()` to reject non-Expr values in contexts where coercion is not wanted; the pythonic coercion is the opposite pattern and should be visually distinct.
276+
**For optional nullable parameters** use `coerce_to_expr_or_none`:
277+
278+
```python
279+
from datafusion.expr import coerce_to_expr, coerce_to_expr_or_none
280+
281+
def regexp_count(
282+
string: Expr,
283+
pattern: Expr | str,
284+
start: Expr | int | None = None,
285+
flags: Expr | str | None = None,
286+
) -> Expr:
287+
pattern = coerce_to_expr(pattern)
288+
start = coerce_to_expr_or_none(start)
289+
flags = coerce_to_expr_or_none(flags)
290+
return Expr(
291+
f.regexp_count(
292+
string.expr,
293+
pattern.expr,
294+
start.expr if start is not None else None,
295+
flags.expr if flags is not None else None,
296+
)
297+
)
298+
```
272299

273-
**Functions with many optional nullable parameters:** For parameters typed as `Expr | <type> | None`, combine the `None` check with the `isinstance` check inline. Repeat this for each parameter — do not factor it into a local helper function, even if the repetition feels verbose. Consistency across the file is more important than DRY within a single function.
300+
Both helpers are defined in `python/datafusion/expr.py` alongside `ensure_expr`. Import them in `functions.py` via:
274301

275302
```python
276-
# For each optional parameter:
277-
if start is not None and not isinstance(start, Expr):
278-
start = Expr.literal(start)
279-
280-
# When passing to the Rust binding, extract .expr or pass None:
281-
f.some_func(
282-
values.expr,
283-
start.expr if start is not None else None,
284-
n.expr if n is not None else None,
285-
)
303+
from datafusion.expr import coerce_to_expr, coerce_to_expr_or_none
286304
```
287305

288306
## What NOT to Change

python/datafusion/expr.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,8 @@
221221
"WindowExpr",
222222
"WindowFrame",
223223
"WindowFrameBound",
224+
"coerce_to_expr",
225+
"coerce_to_expr_or_none",
224226
"ensure_expr",
225227
"ensure_expr_list",
226228
]
@@ -233,6 +235,10 @@ def ensure_expr(value: Expr | Any) -> expr_internal.Expr:
233235
higher level APIs consistently require explicit :func:`~datafusion.col` or
234236
:func:`~datafusion.lit` expressions.
235237
238+
See Also:
239+
:func:`coerce_to_expr` — the opposite behavior: *wraps* non-``Expr``
240+
values as literals instead of rejecting them.
241+
236242
Args:
237243
value: Candidate expression or other object.
238244
@@ -277,6 +283,41 @@ def _iter(
277283
return list(_iter(exprs))
278284

279285

286+
def coerce_to_expr(value: Any) -> Expr:
287+
"""Coerce a native Python value to an ``Expr`` literal, passing ``Expr`` through.
288+
289+
This is the complement of :func:`ensure_expr`: where ``ensure_expr``
290+
*rejects* non-``Expr`` values, ``coerce_to_expr`` *wraps* them via
291+
:meth:`Expr.literal` so that functions can accept native Python types
292+
(``int``, ``float``, ``str``, ``bool``, etc.) alongside ``Expr``.
293+
294+
Args:
295+
value: An ``Expr`` instance (returned as-is) or a Python literal to wrap.
296+
297+
Returns:
298+
An ``Expr`` representing the value.
299+
"""
300+
if isinstance(value, Expr):
301+
return value
302+
return Expr.literal(value)
303+
304+
305+
def coerce_to_expr_or_none(value: Any | None) -> Expr | None:
306+
"""Coerce a value to ``Expr`` or pass ``None`` through unchanged.
307+
308+
Same as :func:`coerce_to_expr` but accepts ``None`` for optional parameters.
309+
310+
Args:
311+
value: An ``Expr`` instance, a Python literal to wrap, or ``None``.
312+
313+
Returns:
314+
An ``Expr`` representing the value, or ``None``.
315+
"""
316+
if value is None:
317+
return None
318+
return coerce_to_expr(value)
319+
320+
280321
def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
281322
"""Convert a Python expression or column name to its raw variant.
282323

0 commit comments

Comments
 (0)