set_assay accepts either an index or assay name (#94)

Author: jkanche

CHANGELOG.md

@@ -1,9 +1,10 @@
 # Changelog
 
-## Version 0.5.4
+## Version 0.5.4 - 0.5.5
 
 - Coercions from SE to RSE and vice-versa.
 - Adapting changes from genomic ranges's search and overlap methods.
+- Update `set_assay` to accept either an assay name or an index position of the assay to replace.
 
 ## Version 0.5.1 - 0.5.3
 

src/summarizedexperiment/BaseSE.py

@@ -960,13 +960,16 @@ def assay(self, assay: Union[int, str]) -> Any:
         """Alias for :py:attr:`~assay`. For backwards compatibility"""
         return self.get_assay(assay)
 
-    def set_assay(self, name: str, assay: Any, in_place: bool = False) -> "BaseSE":
+    def set_assay(self, name: Union[str, int], assay: Any, in_place: bool = False) -> "BaseSE":
         """Add or replace :py:attr:`~summarizedexperiment.BaseSE.BaseSE.assays`'s.
 
         Args:
             name:
                 New or existing assay name.
 
+                Alternatively, may provide an index position of the assay
+                to replace.
+
             assay:
                 A 2-dimensional matrix represented as either
                 :py:class:`~numpy.ndarray` or :py:class:`~scipy.sparse.spmatrix`.
@@ -991,7 +994,19 @@ def set_assay(self, name: str, assay: Any, in_place: bool = False) -> "BaseSE":
         output = self._define_output(in_place)
         if in_place is False:
             output._assays = output._assays.copy()
-        output._assays[name] = assay
+
+        if isinstance(name, int):
+            if name > len(output._assays):
+                raise IndexError("'name' is greather than the number of assays.")
+
+            if name < 0:
+                raise ValueError("'name' cannot be less than 0.")
+
+            output._assays[output.get_assay_names()[name]] = assay
+        elif isinstance(name, str):
+            output._assays[name] = assay
+        else:
+            raise ValueError("'name' must be either a string or an index value.")
         return output
 
     ##########################

tests/test_SE_methods.py

@@ -128,6 +128,17 @@ def test_SE_assay_getters_and_setters():
     assert tse.get_assay("new_counts") is not None
     assert new_tse.get_assay("new_counts") is not None
 
+    mod_tse = tse.set_assay(0, assay=np.random.rand(nrows, ncols), in_place=False)
+    assert mod_tse.get_assay_names() == tse.get_assay_names()
+    assert tse.get_assay("new_counts") is not None
+    assert mod_tse.get_assay("new_counts") is not None
+
+    with pytest.raises(Exception):
+        tse.set_assay(4, assay=np.random.rand(nrows, ncols), in_place=False)
+
+    with pytest.raises(Exception):
+        tse.set_assay(-1, assay=np.random.rand(nrows, ncols), in_place=False)
+
 def test_SE_to_rse():
     tse = SummarizedExperiment(
         assays={"counts": counts}, row_data=row_data, column_data=col_data