feat: autopolars support (#170)

A new feature of jupysql 0.7

pyproject.toml

@@ -26,7 +26,7 @@ python = "^3.8"
 
 duckdb-engine = "^0.7"
 ipython = "^8"
-jupysql = ">=0.6"
+jupysql = ">=0.7"
 pandas = ">=1.5"
 prql-python = "^0.8"
 traitlets = "^5"
@@ -54,6 +54,7 @@ ruff = "^0.0.225"
 safety = "^2"
 sphinx-rtd-theme = "^1.0.0"
 xdoctest = "^0.15.10"
+polars = ">=0.16"
 
 [tool.poetry.scripts]
 pyprql = "pyprql.cli.__init__:main"

pyprql/magic/README.md

@@ -159,6 +159,9 @@ PrqlMagic.autolimit=<Int>
 PrqlMagic.autopandas=<Bool>
     Return Pandas DataFrames instead of regular result sets
     Current: True
+PrqlMagic.autopolars=<Bool>
+    Return Polars DataFrames instead of regular result sets
+    Current: False
 PrqlMagic.autoview=<Bool>
     Display results
     Current: True
@@ -183,6 +186,10 @@ PrqlMagic.dsn_filename=<Unicode>
 PrqlMagic.feedback=<Bool>
     Print number of rows affected by DML
     Current: False
+PrqlMagic.polars_dataframe_kwargs=<key-1>=<value-1>...
+    Polars DataFrame constructor keyword arguments(e.g. infer_schema_length,
+    nan_to_null, schema_overrides, etc)
+    Current: {}
 PrqlMagic.short_errors=<Bool>
     Don't display the full traceback on SQL Programming Error
     Current: True

pyprql/magic/prql.py

@@ -39,6 +39,11 @@ class PrqlMagic(SqlMagic):
         config=True,
         help="Return Pandas DataFrames instead of regular result sets",
     )
+    autopolars = Bool(
+        False,
+        config=True,
+        help="Return Polars DataFrames instead of regular result sets",
+    )
     autoview = Bool(True, config=True, help="Display results")
     feedback = Bool(False, config=True, help="Print number of rows affected by DML")
     target = Unicode("sql.any", config=True, help="Compile target of prql-compiler")

pyprql/tests/test_magic.py

@@ -14,6 +14,7 @@
 import tempfile
 from textwrap import dedent
 
+import polars as pl
 import pytest
 from sqlalchemy import create_engine
 
@@ -262,13 +263,22 @@ def test_bind_vars(ip):
 
 
 def test_autopandas(ip):
-    ip.run_line_magic("config", "SqlMagic.autopandas = True")
     dframe = run_prql(ip, "from test")
     assert not dframe.empty
     assert dframe.ndim == 2
     assert dframe.name[0] == "foo"
 
 
+def test_autopolars(ip):
+    ip.run_line_magic("config", "PrqlMagic.autopolars = True")
+    dframe = run_prql(ip, "from test")
+
+    assert type(dframe) == pl.DataFrame
+    assert not dframe.is_empty()
+    assert len(dframe.shape) == 2
+    assert dframe["name"][0] == "foo"
+
+
 def test_target_dialect(ip):
     ip.run_line_magic("config", 'PrqlMagic.target = "sql.sqlite"')
     dframe = run_prql(