From 39522ccb3f2e76c764cee11e61cdb8b7114e48cc Mon Sep 17 00:00:00 2001 From: David Rodriguez Date: Tue, 11 Jan 2022 16:17:05 -0500 Subject: [PATCH] fix pandas column names; closes 39 --- astrodbkit2/astrodb.py | 2 +- astrodbkit2/tests/test_astrodb.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/astrodbkit2/astrodb.py b/astrodbkit2/astrodb.py index 2d7561c..2258f26 100644 --- a/astrodbkit2/astrodb.py +++ b/astrodbkit2/astrodb.py @@ -332,7 +332,7 @@ def _handle_format(temp, fmt): else: results = AstropyTable(temp) elif fmt.lower() == 'pandas': - results = pd.DataFrame(temp) + results = pd.DataFrame(temp, columns=temp[0].keys()) else: results = temp diff --git a/astrodbkit2/tests/test_astrodb.py b/astrodbkit2/tests/test_astrodb.py index ee93654..45232ec 100644 --- a/astrodbkit2/tests/test_astrodb.py +++ b/astrodbkit2/tests/test_astrodb.py @@ -169,6 +169,11 @@ def test_search_object(mock_simbad, db): t = db.search_object('engu', fuzzy_search=False) assert len(t) == 0 + # Test pandas conversion + t = db.search_object('engu', fmt='pandas') + assert isinstance(t, pd.DataFrame) + assert 'source' in t.columns # check column names + # Search but only consider the Sources.source column t = db.search_object('penguin', table_names={'Sources': ['source']}) assert len(t) == 0