Repository: astropenguin/pandas-dataclasses

Issue:
Fix function for deriving columns

Discussion:

Patch:
diff --git a/pandas_dataclasses/core/api.py b/pandas_dataclasses/core/api.py
index 33753cf..265b988 100644
--- a/pandas_dataclasses/core/api.py
+++ b/pandas_dataclasses/core/api.py
@@ -11,7 +11,7 @@ import numpy as np
 import pandas as pd
 from pandas.api.types import is_list_like
 from typing_extensions import get_origin
-from .specs import Field, Spec
+from .specs import Field, Fields, Spec
 from .tagging import Tag
 from .typing import DataClass, DataClassOf, PAny, TFrame, TPandas, TSeries
 
@@ -132,7 +132,7 @@ def asframe(obj: Any, *, factory: Any = None) -> Any:
     )
 
     dataframe.attrs.update(get_attrs(spec))
-    return dataframe
+    return squeeze(dataframe)
 
 
 @overload
@@ -190,7 +190,7 @@ def asseries(obj: Any, *, factory: Any = None) -> Any:
         series = factory(data=data, index=index, name=name)
 
     series.attrs.update(get_attrs(spec))
-    return series
+    return squeeze(series)
 
 
 def ensure(data: Any, dtype: Optional[str]) -> Any:
@@ -214,16 +214,17 @@ def get_attrs(spec: Spec) -> Dict[Hashable, Any]:
     return data
 
 
-def get_columns(spec: Spec) -> Optional[pd.Index]:
+def get_columns(spec: Spec) -> Optional[pd.MultiIndex]:
     """Derive columns from a specification."""
     if not (fields := spec.fields.of(Tag.DATA)):
         return None
 
-    return squeeze(
-        pd.MultiIndex.from_tuples(
-            map(name, fields),
-            names=fields.names,
-        )
+    if (names := name(fields)) is None:
+        return None
+
+    return pd.MultiIndex.from_tuples(
+        map(name, fields),
+        names=names,
     )
 
 
@@ -238,7 +239,7 @@ def get_data(spec: Spec) -> Dict[Hashable, Any]:
     return data
 
 
-def get_index(spec: Spec) -> Optional[pd.Index]:
+def get_index(spec: Spec) -> Optional[pd.MultiIndex]:
     """Derive index from a specification."""
     if not (fields := spec.fields.of(Tag.INDEX)):
         return None
@@ -249,11 +250,9 @@ def get_index(spec: Spec) -> Optional[pd.Index]:
         for key, val in items(field):
             data[key] = ensure(val, field.dtype)
 
-    return squeeze(
-        pd.MultiIndex.from_arrays(
-            np.broadcast_arrays(*data.values()),
-            names=data.keys(),
-        )
+    return pd.MultiIndex.from_arrays(
+        np.broadcast_arrays(*data.values()),
+        names=data.keys(),
     )
 
 
@@ -265,17 +264,39 @@ def items(field: Field) -> Iterable[Tuple[Hashable, Any]]:
         yield (name(field), field.default)
 
 
-def name(field: Field) -> Hashable:
-    """Derive the name of a field specification."""
-    if isinstance(name := field.name, dict):
-        return tuple(name.values())
-    else:
-        return name
+@overload
+def name(fields: Field) -> Hashable:
+    ...
 
 
-def squeeze(index: pd.Index) -> pd.Index:
-    """Convert a MultiIndex to an Index if possible."""
-    if index.nlevels == 1:
-        return index.get_level_values(0)
-    else:
-        return index
+@overload
+def name(fields: Fields) -> Optional[Hashable]:
+    ...
+
+
+def name(fields: Any) -> Any:
+    """Derive name of a field(s) specification."""
+    if isinstance(fields, Field):
+        if isinstance(name := fields.name, dict):
+            return tuple(name.values())
+        else:
+            return name
+
+    if isinstance(fields, Fields):
+        for field in fields:
+            if isinstance(name := field.name, dict):
+                return tuple(name.keys())
+
+
+def squeeze(data: TPandas) -> TPandas:
+    """Drop levels of an index and columns if possible."""
+    if data.index.nlevels == 1:
+        data.index = data.index.get_level_values(0)
+
+    if isinstance(data, pd.Series):
+        return data  # type: ignore
+
+    if data.columns.nlevels == 1:
+        data.columns = data.columns.get_level_values(0)
+
+    return data
diff --git a/pandas_dataclasses/core/specs.py b/pandas_dataclasses/core/specs.py
index 41684e4..11f8c8f 100644
--- a/pandas_dataclasses/core/specs.py
+++ b/pandas_dataclasses/core/specs.py
@@ -59,13 +59,6 @@ class Field:
 class Fields(Tuple[Field, ...]):
     """List of field specifications with selectors."""
 
-    @property
-    def names(self) -> Optional[Tuple[Hashable, ...]]:
-        """Optional names of the field specifications."""
-        for field in self:
-            if isinstance(name := field.name, dict):
-                return tuple(name.keys())
-
     def of(self, tag: Tag) -> "Fields":
         """Select only fields that have a tag."""
         return type(self)(filter(lambda field: field.has(tag), self))


Test patch:
diff --git a/tests/test_core_api.py b/tests/test_core_api.py
index 31e4ad8..e3f6ef7 100644
--- a/tests/test_core_api.py
+++ b/tests/test_core_api.py
@@ -43,7 +43,7 @@ def test_get_columns() -> None:
     for i in range(len(columns)):
         assert columns[i] == name(spec.fields.of(Tag.DATA)[i])
 
-    assert columns.names == spec.fields.of(Tag.DATA).names  # type: ignore
+    assert columns.names == name(spec.fields.of(Tag.DATA))  # type: ignore
 
 
 def test_get_data() -> None: