get rid of _classify_attr_columns, sprinkle more comments throughout

_pull_attr/_push_attr

Author: ilia-kats

src/mudata/_core/mudata.py

@@ -28,7 +28,6 @@
 from .repr import MUDATA_CSS, block_matrix, details_block_table
 from .utils import (
     MetadataColumn,
-    _classify_attr_columns,
     _make_index_unique,
     _maybe_coerce_to_bool,
     _maybe_coerce_to_boolean,
@@ -1923,9 +1922,29 @@ def _pull_attr(
         if only_drop:
             drop = True
 
-        cols = _classify_attr_columns(
-            {modname: getattr(mod, attr).columns for modname, mod in self.mod.items()}
-        )
+        cols: dict[str, list[MetadataColumn]] = {}
+
+        # get all columns from all modalities and count how many times each column is present
+        derived_name_counts = Counter()
+        for prefix, mod in self.mod.items():
+            modcols = getattr(mod, attr).columns
+            ccols = []
+            for name in modcols:
+                ccols.append(
+                    MetadataColumn(
+                        allowed_prefixes=self.mod.keys(),
+                        prefix=prefix,
+                        name=name,
+                        strip_prefix=False,
+                    )
+                )
+                derived_name_counts[name] += 1
+            cols[prefix] = ccols
+
+        for prefix, modcols in cols.items():
+            for col in modcols:
+                count = derived_name_counts[col.derived_name]
+                col.count = count  # this is important to classify columns
 
         if columns is not None:
             for k, v in {"common": common, "nonunique": nonunique, "unique": unique}.items():
@@ -1936,8 +1955,7 @@ def _pull_attr(
                         stacklevel=2,
                     )
 
-            # - modname1:column -> [modname1:column]
-            # - column -> [modname1:column, modname2:column, ...]
+            # keep only requested columns
             cols = {
                 prefix: [
                     col for col in modcols if col.name in columns or col.derived_name in columns
@@ -1956,15 +1974,18 @@ def _pull_attr(
             if unique is None:
                 unique = True
 
+            # filter columns by class, keep only those that were requested
             selector = {"common": common, "nonunique": nonunique, "unique": unique}
             cols = {
                 prefix: [col for col in modcols if selector[col.klass]]
                 for prefix, modcols in cols.items()
             }
 
+        # filter columns, keep only requested modalities
         if mods is not None:
             cols = {prefix: cols[prefix] for prefix in mods}
 
+        # count final filtered column names, required later to decide whether to prefix a column with its source modality
         derived_name_count = Counter(
             [col.derived_name for modcols in cols.values() for col in modcols]
         )
@@ -2008,6 +2029,8 @@ def _pull_attr(
             if drop:
                 getattr(mod, attr).drop(columns=mod_df.columns, inplace=True)
 
+            # prepend modality prefix to column names if requested via arguments and there are no skipped modalities with
+            # the same column name (prefixing those columns may cause problems with future pulls or pushes)
             mod_df.rename(
                 columns={
                     col.derived_name: col.name
@@ -2027,6 +2050,7 @@ def _pull_attr(
                 inplace=True,
             )
 
+            # reorder modality DF to conform to global order
             mod_df = (
                 _maybe_coerce_to_boolean(mod_df)
                 .iloc[mod_map[mask] - 1]
@@ -2242,6 +2266,7 @@ def _push_attr(
         if only_drop:
             drop = True
 
+        # get all global columns
         cols = [
             MetadataColumn(allowed_prefixes=self.mod.keys(), name=name)
             for name in getattr(self, attr).columns
@@ -2256,9 +2281,7 @@ def _push_attr(
                         stacklevel=2,
                     )
 
-            # - modname1:column -> [modname1:column]
-            # - column -> [modname1:column, modname2:column, ...]
-            # preemptively drop columns from other modalities
+            # keep only requested columns
             cols = [
                 col
                 for col in cols
@@ -2271,8 +2294,8 @@ def _push_attr(
             if prefixed is None:
                 prefixed = True
 
+            # filter columns by class, keep only those that were requested
             selector = {"common": common, "unknown": prefixed}
-
             cols = [col for col in cols if selector[col.klass]]
 
         if len(cols) == 0:
@@ -2290,20 +2313,22 @@ def _push_attr(
                 )
 
         attrmap = getattr(self, f"{attr}map")
-        _n_attr = self.n_vars if attr == "var" else self.n_obs
-
         for m, mod in self.mod.items():
             if mods is not None and m not in mods:
                 continue
 
             mod_map = attrmap[m].ravel()
-            mask = mod_map != 0
-            mod_n_attr = mod.n_vars if attr == "var" else mod.n_obs
+            mask = mod_map > 0
+            mod_n_attr = mod.n_obs if attr == "obs" else mod.n_vars
 
+            # get all common and modality-specific columns for the current modality
             mod_cols = [col for col in cols if col.prefix == m or col.klass == "common"]
             df = getattr(self, attr)[mask].loc[:, [col.name for col in mod_cols]]
+
+            # strip modality prefix where necessary
             df.columns = [col.derived_name for col in mod_cols]
 
+            # reorder global DF to conform to modality order
             df = df.iloc[np.argsort(mod_map[mask])].set_index(np.arange(mod_n_attr))
 
             if not only_drop:

src/mudata/_core/utils.py

@@ -1,5 +1,5 @@
 from collections import Counter
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 from typing import Literal, TypeVar
 
 import numpy as np
@@ -39,7 +39,7 @@ def _maybe_coerce_to_boolean(df: T) -> T:
 
 
 class MetadataColumn:
-    __slots__ = ("prefix", "derived_name", "count", "_allowed_prefixes")
+    __slots__ = ("prefix", "derived_name", "count", "_allowed_prefixes", "_strip_prefix")
 
     def __init__(
         self,
@@ -48,9 +48,12 @@ def __init__(
         prefix: str | None = None,
         name: str | None = None,
         count: int = 0,
+        strip_prefix: bool = True,
     ):
+        self._strip_prefix = strip_prefix
         self._allowed_prefixes = allowed_prefixes
-        if prefix is None:
+        self.prefix = prefix
+        if prefix is None and strip_prefix:
             self.name = name
         else:
             self.prefix = prefix
@@ -67,7 +70,8 @@ def name(self) -> str:
     @name.setter
     def name(self, new_name):
         if (
-            len(name_split := new_name.split(":", 1)) < 2
+            not self._strip_prefix
+            or len(name_split := new_name.split(":", 1)) < 2
             or name_split[0] not in self._allowed_prefixes
         ):
             self.prefix = None
@@ -87,37 +91,6 @@ def klass(self) -> Literal["common", "unique", "nonunique", "unknown"]:
             return "unknown"
 
 
-def _classify_attr_columns(names: Mapping[str, Sequence[str]]) -> dict[str, list[MetadataColumn]]:
-    """
-    Classify names into common, non-unique, and unique
-    w.r.t. to the list of prefixes.
-
-    - Common columns do not have modality prefixes.
-    - Non-unqiue columns have a modality prefix,
-      and there are multiple columns that differ
-      only by their modality prefix.
-    - Unique columns are prefixed by modality names,
-      and there is only one modality prefix
-      for a column with a certain name.
-    """
-    res: dict[str, list[MetadataColumn]] = {}
-
-    derived_name_counts = Counter()
-    for prefix, pnames in names.items():
-        cres = []
-        for name in pnames:
-            cres.append(MetadataColumn(allowed_prefixes=names.keys(), prefix=prefix, name=name))
-            derived_name_counts[name] += 1
-        res[prefix] = cres
-
-    for prefix, names in res.items():
-        for name_res in names:
-            count = derived_name_counts[name_res.derived_name]
-            name_res.count = count
-
-    return res
-
-
 def _update_and_concat(df1: pd.DataFrame, df2: pd.DataFrame) -> pd.DataFrame:
     df = df1.copy(deep=False)
     # This converts boolean to object dtype, unfortunately