(refactor): fix_known_differences usage of as_type

scverse · Aug 27, 2024 · b1feb6f · b1feb6f
1 parent 2cf1262
commit b1feb6f
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/tests/test_concatenate.py b/tests/test_concatenate.py
@@ -149,12 +149,11 @@ def fix_known_differences(orig, result, backwards_compat=True):
         result.obs.drop(columns=["batch"], inplace=True)
 
     # Possibly need to fix this, ordered categoricals lose orderedness
-    for k, dtype in orig.obs.dtypes.items():
-        if isinstance(dtype, pd.CategoricalDtype) and dtype.ordered:
-            result.obs[k] = result.obs[k].astype(dtype)
-    for k, dtype in orig.obsm["df"].dtypes.items():
-        if isinstance(dtype, pd.CategoricalDtype) and dtype.ordered:
-            result.obsm["df"][k] = result.obsm["df"][k].astype(dtype)
+    for get_df in [lambda k: getattr(k, "obs"), lambda k: getattr(k, "obsm")["df"]]:
+        str_to_df_converted = get_df(result)
+        for k, dtype in get_df(orig).dtypes.items():
+            if isinstance(dtype, pd.CategoricalDtype) and dtype.ordered:
+                str_to_df_converted[k] = str_to_df_converted[k].astype(dtype)
 
     return orig, result