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