From 52c0212c04ee7ce8921cfccbb6e86d9333b47a50 Mon Sep 17 00:00:00 2001 From: Nathan McDougall Date: Sat, 20 Jul 2024 06:58:23 +1200 Subject: [PATCH] Support GeoDataFrame in `default_title` --- pins/drivers.py | 12 +++++++++++- pins/tests/test_drivers.py | 5 +++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/pins/drivers.py b/pins/drivers.py index b9a7c72..2cc5e79 100644 --- a/pins/drivers.py +++ b/pins/drivers.py @@ -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" diff --git a/pins/tests/test_drivers.py b/pins/tests/test_drivers.py index f6d3385..71d4b44 100644 --- a/pins/tests/test_drivers.py +++ b/pins/tests/test_drivers.py @@ -1,6 +1,7 @@ from pathlib import Path import fsspec +import geopandas as gpd import pandas as pd import pytest @@ -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"),