Cleanup and add asserters to grep_annotation_column

src/base/openpipelinetestutils/utils.py

@@ -13,7 +13,7 @@ def remove_annotation_column(annotation_object: AnnotationObject,
     if isinstance(annotation_object, AnnData) and modality_name is not None:
         raise ValueError("Cannot specify modality when object is of type AnnData.")
     if isinstance(column_names, str):
-        column_names = [column_names]
+        column_names = [str(column_names)] # str to make a copy
     axis_strings = {
         "var": "var",
         "obs": "obs",
@@ -23,9 +23,10 @@ def remove_annotation_column(annotation_object: AnnotationObject,
     axis_string = axis_strings[axis]
     axis_getter = attrgetter(axis_string)
     axis_setter = lambda obj, value: setattr(obj, axis_string, value)
-    axis_setter(annotation_object, axis_getter(annotation_object).drop(column_names,
-                                                                       axis="columns",
-                                                                       inplace=False))
+    if not modality_name:
+        axis_setter(annotation_object, axis_getter(annotation_object).drop(column_names,
+                                                                        axis="columns",
+                                                                        inplace=False))
     if isinstance(annotation_object, MuData):
         modality_names = [modality_name] if modality_name else list(annotation_object.mod.keys())
         extra_cols_to_remove = [f"{mod_name}:{column_name}" for mod_name, column_name

src/metadata/add_id/test.py

@@ -15,18 +15,6 @@
 }
 ## VIASH END
 
-
-# def remove_obs_column(mudata_object, column):
-#     for _, modality in mudata_object.mod.items():
-#         modality.obs = modality.obs.drop([column],
-#                                          axis="columns", inplace=False)
-#     mod_names_columns = [f"{mod_name}:{column}" for mod_name in mudata_object.mod.keys()]
-#     mudata_object.obs = mudata_object.obs.drop([column] + mod_names_columns,
-#                                                axis="columns",
-#                                                inplace=False)
-#     mudata_object.update()
-#     return mudata_object
-
 @pytest.fixture
 def generate_h5mu():
     # generate data

src/metadata/grep_annotation_column/test.py

@@ -7,7 +7,9 @@
 from anndata import AnnData
 from mudata import MuData, read_h5mu
 from subprocess import CalledProcessError
-from uuid import uuid4
+from openpipelinetestutils.asserters import assert_annotation_objects_equal
+from openpipelinetestutils.utils import remove_annotation_column
+
 
 ## VIASH START
 meta = {
@@ -18,7 +20,6 @@
 }
 ## VIASH END
 
-from openpipelinetestutils.asserters import assert_annotation_objects_equal
 
 @pytest.fixture
 def generate_h5mu():
@@ -34,21 +35,6 @@ def generate_h5mu():
     tmp_mudata = MuData({'mod1': ad1, 'mod2': ad2})
     return tmp_mudata
 
-@pytest.fixture
-def random_h5mu_path(tmp_path):
-    def wrapper():
-        return tmp_path / f"{uuid4()}.h5mu"
-    return wrapper
-
-
-@pytest.fixture
-def write_mudata_to_file(random_h5mu_path):
-    def wrapper(mudata_obj):
-        output_path = random_h5mu_path()
-        mudata_obj.write(output_path)
-        return output_path
-    return wrapper
-
 
 @pytest.mark.parametrize("compression_format", ["gzip", "lzf"])
 def test_grep_column(run_component, generate_h5mu,
@@ -75,6 +61,10 @@ def test_grep_column(run_component, generate_h5mu,
     assert "test" in output_data.mod['mod1'].var.columns.to_list()
     assert output_data.mod['mod1'].var['test'].to_list() == [True, False, False]
 
+    assert_annotation_objects_equal(input_path,
+                                    remove_annotation_column(output_data, "test", "var", "mod1"),
+                                    check_data=True)
+
 @pytest.mark.parametrize("compression_format", ["gzip", "lzf"])
 def test_grep_column_default(run_component, generate_h5mu,
                              random_h5mu_path, write_mudata_to_file,
@@ -99,6 +89,10 @@ def test_grep_column_default(run_component, generate_h5mu,
     assert "test" in output_data.mod['mod1'].var.columns.to_list()
     assert output_data.mod['mod1'].var['test'].to_list() == [True, False, False]
 
+    assert_annotation_objects_equal(input_path,
+                                    remove_annotation_column(output_data, "test", "var", "mod1"),
+                                    check_data=True)
+
 @pytest.mark.parametrize("compression_format", ["gzip", "lzf"])
 def test_grep_column(run_component, generate_h5mu,
                      random_h5mu_path, write_mudata_to_file,
@@ -124,6 +118,9 @@ def test_grep_column(run_component, generate_h5mu,
     assert "test" in output_data.mod['mod1'].var.columns.to_list()
     assert output_data.mod['mod1'].var['test'].to_list() == [True, False, False]
 
+    assert_annotation_objects_equal(input_path,
+                                    remove_annotation_column(output_data, "test", "var", "mod1"),
+                                    check_data=True)
 
 @pytest.mark.parametrize("compression_format", ["gzip", "lzf"])
 def test_grep_column_fraction_column(run_component, generate_h5mu,
@@ -151,7 +148,15 @@ def test_grep_column_fraction_column(run_component, generate_h5mu,
     assert "test" in output_data.mod['mod1'].var.columns.to_list()
     assert output_data.mod['mod1'].var['test'].to_list() == [True, False, False]
     output_data.mod['mod1'].obs['test_output_fraction'].to_list() == [1/6, 4/15]
-
+    output_object_without_obs_column = remove_annotation_column(output_data, 
+                                                                ["test_output_fraction"], 
+                                                                "obs", "mod1")
+    output_object_without_wo_output = remove_annotation_column(output_object_without_obs_column, 
+                                                               ["test"], 
+                                                               "var", "mod1")
+    assert_annotation_objects_equal(input_path,
+                                    output_object_without_wo_output,
+                                    check_data=True)
 
 @pytest.mark.parametrize("compression_format", ["gzip", "lzf"])
 def test_fraction_column_input_layer(run_component, generate_h5mu,
@@ -182,6 +187,16 @@ def test_fraction_column_input_layer(run_component, generate_h5mu,
     assert "test" in output_data.mod['mod1'].var.columns.to_list()
     assert output_data.mod['mod1'].var['test'].to_list() == [True, False, False]
     output_data.mod['mod1'].obs['test_output_fraction'].to_list() == [1/6, 4/15]
+    output_object_without_obs_column = remove_annotation_column(output_data, 
+                                                                ["test_output_fraction"], 
+                                                                "obs", "mod1")
+    output_object_without_wo_output = remove_annotation_column(output_object_without_obs_column, 
+                                                               ["test"], 
+                                                               "var", "mod1")
+    assert_annotation_objects_equal(input_path,
+                                    output_object_without_wo_output,
+                                    check_data=True)
+
 
 def test_missing_column(run_component, generate_h5mu,
                         random_h5mu_path, write_mudata_to_file):