From b1feb6f3b060a7691eb1d3d507caa71579ad4ef6 Mon Sep 17 00:00:00 2001 From: ilan-gold Date: Tue, 27 Aug 2024 16:48:44 +0200 Subject: [PATCH] (refactor): `fix_known_differences` usage of `as_type` --- tests/test_concatenate.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/test_concatenate.py b/tests/test_concatenate.py index d7b3f4767..607b5b613 100644 --- 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