diff --git a/leanframe/core/frame.py b/leanframe/core/frame.py index 3010e43..214c7d3 100644 --- a/leanframe/core/frame.py +++ b/leanframe/core/frame.py @@ -30,5 +30,10 @@ def __init__(self, data): else: raise NotImplementedError("DataFrame constructor doesn't support local data yet.") + @property + def columns(self) -> pandas.Index: + """The column labels of the DataFrame.""" + return pandas.Index(self._data.columns, dtype="object") + def to_pandas(self) -> pandas.DataFrame: return self._data.to_pandas() diff --git a/tests/unit/session/test_read_sql_table.py b/tests/unit/session/test_read_sql_table.py index affe86a..2e8e198 100644 --- a/tests/unit/session/test_read_sql_table.py +++ b/tests/unit/session/test_read_sql_table.py @@ -16,6 +16,8 @@ import leanframe +import pandas.testing + def test_read_sql_table_can_convert_to_pandas(session: leanframe.Session): """Read a table with simple scalar values.""" @@ -26,3 +28,6 @@ def test_read_sql_table_can_convert_to_pandas(session: leanframe.Session): # Make sure we have _some_ data. assert len(df_pd.index) > 0 assert len(df_pd.columns) > 0 + + # Where possible, check that we are compatible with pandas. + pandas.testing.assert_index_equal(df_pd.columns, df_lf.columns)