Support polars.DataFrame in default_title

Author: nathanjmcdougall

pins/drivers.py

@@ -233,13 +233,15 @@ def save_data(obj, fname, type=None, apply_suffix: bool = True) -> "str | Sequen
 
 
 def default_title(obj, name):
-    import pandas as pd
+    df_family = _get_df_family(obj)
 
-    if isinstance(obj, pd.DataFrame):
+    if df_family in ("pandas", "polars"):
         # TODO(compat): title says CSV rather than data.frame
         # see https://github.com/machow/pins-python/issues/5
         shape_str = " x ".join(map(str, obj.shape))
         return f"{name}: a pinned {shape_str} DataFrame"
-    else:
+    elif df_family == "unknown":
         obj_name = type(obj).__qualname__
         return f"{name}: a pinned {obj_name} object"
+    else:
+        assert_never(df_family)

pins/tests/test_drivers.py

@@ -1,6 +1,7 @@
 import fsspec
 import pytest
 import pandas as pd
+import polars as pl
 
 from pathlib import Path
 
@@ -35,6 +36,7 @@ class D:
     [
         (pd.DataFrame({"x": [1, 2]}), "somename: a pinned 2 x 1 DataFrame"),
         (pd.DataFrame({"x": [1], "y": [2]}), "somename: a pinned 1 x 2 DataFrame"),
+        (pl.DataFrame({"x": [1, 2]}), "somename: a pinned 2 x 1 DataFrame"),
         (ExC(), "somename: a pinned ExC object"),
         (ExC().D(), "somename: a pinned ExC.D object"),
         ([1, 2, 3], "somename: a pinned list object"),