Skip to content

Commit

Permalink
Support GeoDataFrame in default_title
Browse files Browse the repository at this point in the history
  • Loading branch information
nathanjmcdougall committed Jul 25, 2024
1 parent f477398 commit 52c0212
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 1 deletion.
12 changes: 11 additions & 1 deletion pins/drivers.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,10 +231,20 @@ def default_title(obj, name):
import pandas as pd

if isinstance(obj, pd.DataFrame):
try:
import geopandas as gpd
except ModuleNotFoundError:
obj_name = "DataFrame"
else:
if isinstance(obj, gpd.GeoDataFrame):
obj_name = "GeoDataFrame"
else:
obj_name = "DataFrame"

# 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"
return f"{name}: a pinned {shape_str} {obj_name}"
else:
obj_name = type(obj).__qualname__
return f"{name}: a pinned {obj_name} object"
5 changes: 5 additions & 0 deletions pins/tests/test_drivers.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from pathlib import Path

import fsspec
import geopandas as gpd
import pandas as pd
import pytest

Expand Down Expand Up @@ -34,6 +35,10 @@ 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"),
(
gpd.GeoDataFrame({"x": [1], "geometry": [None]}),
"somename: a pinned 1 x 2 GeoDataFrame",
),
(ExC(), "somename: a pinned ExC object"),
(ExC().D(), "somename: a pinned ExC.D object"),
([1, 2, 3], "somename: a pinned list object"),
Expand Down

0 comments on commit 52c0212

Please sign in to comment.