diff --git a/changelog.md b/changelog.md index 89b5d149..3b8d8fef 100644 --- a/changelog.md +++ b/changelog.md @@ -8,6 +8,7 @@ Features * Display of inline plots returned from `.|` operations. * Don't attempt inline plots in the Windows console. * Save plots, like parquets, with the `.>` operator. +* Display many native Python types returned from `.|` operations. Bug Fixes diff --git a/doc/transforms.md b/doc/transforms.md index 52338c22..bf0bc2d5 100644 --- a/doc/transforms.md +++ b/doc/transforms.md @@ -53,9 +53,12 @@ SELECT customer_id, COUNT(1) AS len FROM orders GROUP BY customer_id; ``` Transform expressions run with normal Python privileges, and expressions -should not be run from untrusted sources. If the transform operation -returns a Polars `DataFrame` or `Series`, the result is rendered by mycli -as tabular output. Most other return types will be silently ignored. +should not be run from untrusted sources. + +If the transform operation returns a Polars `DataFrame` or `Series`, or a +Python type which can be rendered as a table, the result is rendered by mycli +as tabular output. `None` return values will be silently ignored, and other +return types will result in a warning message. Transform expressions are useful for operations such as medians which cannot be done (or are awkward) in SQL. Example: diff --git a/mycli/packages/polars_transform.py b/mycli/packages/polars_transform.py index 3b0b6a59..0c8f51ba 100644 --- a/mycli/packages/polars_transform.py +++ b/mycli/packages/polars_transform.py @@ -2,9 +2,12 @@ import builtins from dataclasses import dataclass +import datetime +import decimal +import fractions from io import BytesIO from types import CodeType -from typing import Any, Iterable +from typing import Any, Iterable, Sequence import sqlglot @@ -36,6 +39,25 @@ class PolarsTransform: altair: Any | None +def _is_coerceable_scalarish(value: Any) -> bool: + # pretend str/bytes are scalars + return isinstance( + value, + ( + bool, + int, + float, + complex, + bytes, + str, + datetime.datetime, + datetime.timedelta, + decimal.Decimal, + fractions.Fraction, + ), + ) + + def parse_polars_transform(command: str) -> PolarsPipeline | None: """Parse a SQL statement with optional Polars transform and file output.""" try: @@ -207,10 +229,32 @@ def run_polars_transform( try: value = eval( transform.code, - {'__builtins__': builtins, 'df': dataframe, 'pl': transform.polars, 'alt': transform.altair}, + { + '__builtins__': builtins, + 'df': dataframe, + 'pl': transform.polars, + 'alt': transform.altair, + }, ) except Exception as exc: raise PolarsTransformError(f'Polars expression failed: {type(exc).__name__}: {exc}') from exc + + if _is_coerceable_scalarish(value): + try: + value = transform.polars.DataFrame([value], schema=[str(value)]) + except Exception as exc: + raise PolarsTransformError(f'Unable to render scalar as DataFrame: {type(exc).__name__}: {exc}') from exc + elif isinstance(value, dict): + try: + value = transform.polars.DataFrame(value) + except Exception as exc: + raise PolarsTransformError(f'Unable to render dictionary as DataFrame: {type(exc).__name__}: {exc}') from exc + elif isinstance(value, Sequence): + try: + value = transform.polars.Series(value, strict=False) + except Exception as exc: + raise PolarsTransformError(f'Unable to render Sequence as Series: {type(exc).__name__}: {exc}') from exc + if isinstance(value, transform.polars.DataFrame): if output_path is not None: if not output_path.lower().endswith('.parquet'): @@ -221,7 +265,7 @@ def run_polars_transform( raise PolarsTransformError(f'Unable to write Parquet file "{output_path}": {type(exc).__name__}: {exc}') from exc return SQLResult(status=f'Wrote {len(value)} rows to {output_path}.') return SQLResult(header=list(value.columns), rows=list(value.iter_rows())) - if isinstance(value, transform.polars.Series): + elif isinstance(value, transform.polars.Series): column_name = value.name or 'value' if output_path is not None: if not output_path.lower().endswith('.parquet'): @@ -233,7 +277,7 @@ def run_polars_transform( raise PolarsTransformError(f'Unable to write Parquet file "{output_path}": {type(exc).__name__}: {exc}') from exc return SQLResult(status=f'Wrote {len(series_dataframe)} rows to {output_path}.') return SQLResult(header=[column_name], rows=[(item,) for item in value]) - if transform.altair is not None and isinstance(value, transform.altair.TopLevelMixin): + elif transform.altair is not None and isinstance(value, transform.altair.TopLevelMixin): if output_path is not None and not output_path.lower().endswith('.png'): raise PolarsTransformError('Altair charts can only be written to ".png" files.') if output_path is None and image_protocol == 'none': @@ -265,8 +309,7 @@ def run_polars_transform( except Exception as exc: raise PolarsTransformError(f'Unable to render Altair chart: {type(exc).__name__}: {exc}') from exc return SQLResult(image=png.getvalue(), image_protocol=image_protocol) - if output_path is not None: - raise PolarsTransformError( - 'Polars transforms must return a DataFrame or Series for Parquet output, or an Altair chart for PNG output.' - ) + elif value is None: + return SQLResult() + return SQLResult(status=f'Nothing could be displayed for return type: {type(value)}') diff --git a/test/pytests/test_polars_transform.py b/test/pytests/test_polars_transform.py index 89b0b351..65f1f085 100644 --- a/test/pytests/test_polars_transform.py +++ b/test/pytests/test_polars_transform.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Any, Iterator +from pathlib import Path +from typing import Any, Iterator, Sequence import pytest @@ -21,9 +22,17 @@ class FakeDataFrame: written_paths: list[str] = [] written_dataframes: list['FakeDataFrame'] = [] - def __init__(self, rows: list[tuple[Any, ...]], schema: list[str], orient: str) -> None: - assert orient == 'row' - self.rows = rows + def __init__( + self, + rows: list[tuple[Any, ...]] | list[Any], + schema: list[str], + orient: str | None = None, + ) -> None: + if orient is None: + self.rows = [(value,) for value in rows] + else: + assert orient == 'row' + self.rows = rows self.columns = schema self.parquet_paths: list[str] = [] @@ -40,9 +49,20 @@ def write_parquet(self, path: str) -> None: class FakeSeries: - def __init__(self, name: str, values: list[Any]) -> None: - self.name = name - self.values = values + def __init__( + self, + name: str | Sequence[Any], + values: list[Any] | None = None, + *, + strict: bool = True, + ) -> None: + if values is None: + self.name = '' + self.values = list(name) + else: + assert isinstance(name, str) + self.name = name + self.values = values def __iter__(self) -> Iterator[Any]: return iter(self.values) @@ -68,6 +88,41 @@ class FailingPolars: DataFrame = FailingDataFrame +class ScalarConstructionFailingDataFrame(FakeDataFrame): + def __init__( + self, + rows: list[tuple[Any, ...]] | list[Any], + schema: list[str], + orient: str | None = None, + ) -> None: + if orient is None: + raise TypeError('invalid scalar') + super().__init__(rows, schema, orient) + + +class ScalarConstructionFailingPolars: + DataFrame = ScalarConstructionFailingDataFrame + Series = FakeSeries + + +class SequenceConstructionFailingSeries(FakeSeries): + def __init__( + self, + name: str | Sequence[Any], + values: list[Any] | None = None, + *, + strict: bool = True, + ) -> None: + if values is None: + raise TypeError('invalid sequence') + super().__init__(name, values, strict=strict) + + +class SequenceConstructionFailingPolars: + DataFrame = FakeDataFrame + Series = SequenceConstructionFailingSeries + + class FakeTopLevelMixin: pass @@ -334,6 +389,158 @@ def test_run_polars_transform_materializes_and_renders_returned_dataframe() -> N assert result == SQLResult(header=['id'], rows=[(1,), (2,)]) +@pytest.mark.parametrize( + ('expression', 'column_name', 'value'), + [ + ('42', '42', 42), + ("'total'", 'total', 'total'), + ], +) +def test_run_polars_transform_coerces_scalar_to_dataframe( + expression: str, + column_name: str, + value: Any, +) -> None: + result = run_polars_transform( + make_transform(expression), + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + assert result == SQLResult(header=[column_name], rows=[(value,)]) + + +def test_run_polars_transform_reports_invalid_scalar() -> None: + transform = PolarsTransform( + sql='SELECT id FROM orders', + expression='42', + code=compile('42', '', 'eval'), + polars=ScalarConstructionFailingPolars, + altair=None, + ) + + with pytest.raises(PolarsTransformError, match='Unable to render scalar as DataFrame: TypeError: invalid scalar'): + run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + +@pytest.mark.parametrize( + ('expression', 'header', 'rows'), + [ + ("{'id': [1, 2], 'name': ['one', 'two']}", ['id', 'name'], [(1, 'one'), (2, 'two')]), + ("{'id': 1, 'name': 'one'}", ['id', 'name'], [(1, 'one')]), + ('{}', [], []), + ], +) +def test_run_polars_transform_coerces_dictionary_to_dataframe( + expression: str, + header: list[str], + rows: list[tuple[Any, ...]], +) -> None: + import polars as pl + + transform = PolarsTransform( + sql='SELECT id FROM orders', + expression=expression, + code=compile(expression, '', 'eval'), + polars=pl, + altair=None, + ) + + result = run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + assert result == SQLResult(header=header, rows=rows) + + +def test_run_polars_transform_coerces_sequence_to_series() -> None: + result = run_polars_transform( + make_transform('[1, None, 3]'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + assert result == SQLResult(header=['value'], rows=[(1,), (None,), (3,)]) + + +def test_run_polars_transform_reports_invalid_sequence() -> None: + transform = PolarsTransform( + sql='SELECT id FROM orders', + expression='[1, 2]', + code=compile('[1, 2]', '', 'eval'), + polars=SequenceConstructionFailingPolars, + altair=None, + ) + + with pytest.raises(PolarsTransformError, match='Unable to render Sequence as Series: TypeError: invalid sequence'): + run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + +def test_run_polars_transform_returns_empty_result_for_none() -> None: + result = run_polars_transform( + make_transform('None'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + assert result == SQLResult() + + +def test_run_polars_transform_reports_unsupported_return_type() -> None: + result = run_polars_transform( + make_transform('object()'), + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + assert result.status_plain == "Nothing could be displayed for return type: " + + +def test_run_polars_transform_reports_invalid_dictionary() -> None: + import polars as pl + + expression = "{'id': [1, 2], 'name': ['one']}" + transform = PolarsTransform( + sql='SELECT id FROM orders', + expression=expression, + code=compile(expression, '', 'eval'), + polars=pl, + altair=None, + ) + + with pytest.raises(PolarsTransformError, match='Unable to render dictionary as DataFrame'): + run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + ) + + +def test_run_polars_transform_writes_dictionary_to_parquet(tmp_path: Path) -> None: + import polars as pl + + path = tmp_path / 'orders.parquet' + expression = "{'id': [1, 2], 'name': ['one', 'two']}" + transform = PolarsTransform( + sql='SELECT id FROM orders', + expression=expression, + code=compile(expression, '', 'eval'), + polars=pl, + altair=None, + ) + + result = run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,)])]), + str(path), + ) + + assert result == SQLResult(status=f'Wrote 2 rows to {path}.') + assert pl.read_parquet(path).to_dict(as_series=False) == {'id': [1, 2], 'name': ['one', 'two']} + + @pytest.mark.parametrize( ('expression', 'header', 'rows'), [ @@ -427,24 +634,6 @@ def test_run_polars_transform_reports_parquet_write_error() -> None: ) -def test_run_polars_transform_rejects_non_tabular_parquet_output() -> None: - with pytest.raises(PolarsTransformError, match='must return a DataFrame or Series'): - run_polars_transform( - make_transform('1'), - iter([SQLResult(header=['id'], rows=[(1,)])]), - 'orders.parquet', - ) - - -def test_run_polars_transform_reports_non_dataframe_result() -> None: - result = run_polars_transform( - make_transform('1'), - iter([SQLResult(header=['id'], rows=[(1,)])]), - ) - - assert result.status_plain == "Nothing could be displayed for return type: " - - def test_run_polars_transform_reports_disabled_altair_chart_output() -> None: result = run_polars_transform( make_transform('alt.Chart(df)'),