Repository: astropenguin/pandas-dataclasses

Issue:
Fix wrong returns of get_columns

Discussion:

Patch:
diff --git a/pandas_dataclasses/core/asdata.py b/pandas_dataclasses/core/asdata.py
index 1b31a11..57a0226 100644
--- a/pandas_dataclasses/core/asdata.py
+++ b/pandas_dataclasses/core/asdata.py
@@ -41,8 +41,12 @@ def asdataframe(obj: Any, *, factory: Any = None) -> Any:
     if not issubclass(factory, pd.DataFrame):
         raise TypeError("Factory must be a subclass of DataFrame.")
 
-    dataframe = factory(data=get_data(spec), index=get_index(spec))
-    dataframe.columns.names = get_columns(spec)
+    dataframe = factory(
+        data=get_data(spec),
+        index=get_index(spec),
+        columns=get_columns(spec),
+    )
+
     dataframe.attrs.update(get_attrs(spec))
     return dataframe
 
@@ -104,14 +108,17 @@ def get_attrs(spec: Spec) -> "dict[Hashable, Any]":
     return objs
 
 
-def get_columns(spec: Spec) -> "list[Hashable]":
-    """Derive column names from a specification."""
-    objs: "dict[Hashable, Any]" = {}
-
-    for field in spec.fields.of_column:
-        objs[field.name] = field.default
+def get_columns(spec: Spec) -> Optional[pd.Index]:
+    """Derive columns from a specification."""
+    names = [field.name for field in spec.fields.of_column]
+    elems = [field.name for field in spec.fields.of_data]
 
-    return list(objs) if objs else [None]
+    if len(names) == 0:
+        return
+    if len(names) == 1:
+        return pd.Index(pd.array(elems), name=names[0])
+    else:
+        return pd.MultiIndex.from_tuples(elems, names=names)
 
 
 def get_data(spec: Spec) -> "dict[Hashable, Any]":


Test patch:
diff --git a/tests/test_asdata.py b/tests/test_asdata.py
index 2d0e543..2d40b13 100644
--- a/tests/test_asdata.py
+++ b/tests/test_asdata.py
@@ -68,10 +68,15 @@ def test_get_attrs() -> None:
 
 
 def test_get_columns() -> None:
-    columns = get_columns(spec)
+    columns = cast(pd.Index, get_columns(spec))
 
-    assert columns[0] == spec.fields.of_column[0].name
-    assert columns[1] == spec.fields.of_column[1].name
+    assert columns.names[0] == spec.fields.of_column[0].name
+    assert columns.names[1] == spec.fields.of_column[1].name
+
+    assert columns[0] == spec.fields.of_data[0].name
+    assert columns[1] == spec.fields.of_data[1].name
+    assert columns[2] == spec.fields.of_data[2].name
+    assert columns[3] == spec.fields.of_data[3].name
 
 
 def test_get_data() -> None: