Added generic for assigning values to sequences/rows.

This is basically the equivalent of Bioconductor's replaceROWS generic.

src/biocutils/__init__.py

@@ -41,6 +41,10 @@
 from .subset_rows import subset_rows
 from .subset_sequence import subset_sequence
 
+from .assign import assign
+from .assign_rows import assign_rows
+from .assign_sequence import assign_sequence
+
 from .show_as_cell import show_as_cell
 from .convert_to_dense import convert_to_dense
 

src/biocutils/assign.py

@@ -0,0 +1,25 @@
+from typing import Any, Sequence
+
+from .assign_rows import assign_rows
+from .assign_sequence import assign_sequence
+from .is_high_dimensional import is_high_dimensional
+
+
+def assign(x: Any, indices: Sequence[int], replacement: Any) -> Any:
+    """
+    Generic assign that checks if the objects are n-dimensional for n > 1 (i.e.
+    has a ``shape`` property of length greater than 1); if so, it calls
+    :py:func:`~biocutils.assign_rows.assign_rows` to assign them along the
+    first dimension, otherwise it assumes that they are vector-like and calls
+    :py:func:`~biocutils.assign_sequence.assign_sequence` instead.
+
+    Args:
+        x: Object to be assignted.
+
+    Returns:
+        The object after assignment, typically the same type as ``x``.
+    """
+    if is_high_dimensional(x):
+        return assign_rows(x, indices, replacement)
+    else:
+        return assign_sequence(x, indices, replacement)

src/biocutils/assign_rows.py

@@ -0,0 +1,36 @@
+from typing import Any, Sequence, Union
+from functools import singledispatch
+import numpy
+
+
+@singledispatch
+def assign_rows(x: Any, indices: Sequence[int], replacement: Any) -> Any:
+    """
+    Assign ``replacement`` values to a copy of ``x`` at the rows specified by
+    ``indices``. This 
+
+    Args:
+        x:
+            Any high-dimensional object.
+
+        indices: 
+            Sequence of non-negative integers specifying rows of ``x``.
+
+        replacement:
+            Replacement values to be assigned to ``x``. This should have the
+            same number of rows as the length of ``indices``. Typically
+            ``replacement`` will have the same dimensionality as ``x``.
+
+    Returns:
+        A copy of ``x`` with the rows replaced at ``indices``.
+    """
+    raise NotImplementedError("no `assign_rows` method implemented for '" + type(x[0]).__name__ + "' objects")
+
+
+@assign_rows.register
+def _assign_rows_numpy(x: numpy.ndarray, indices: Sequence[int], replacement: Any) -> numpy.ndarray:
+    tmp = [slice(None)] * len(x.shape)
+    tmp[0] = indices 
+    output = numpy.copy(x)
+    output[(*tmp,)] = replacement
+    return output

src/biocutils/assign_sequence.py

@@ -0,0 +1,51 @@
+from typing import Any, Sequence, Union
+from functools import singledispatch
+import numpy
+
+
+@singledispatch
+def assign_sequence(x: Any, indices: Sequence[int], replacement: Any) -> Any:
+    """
+    Assign ``replacement`` values to a copy of ``x`` at the specified ``indices``.
+
+    Args:
+        x: 
+            Any sequence-like object that can be assigned.
+
+        indices: 
+            Sequence of non-negative integers specifying positions on ``x``.
+
+        replacement:
+            Replacement values to be assigned to ``x``. This should have the
+            same length as ``indices``.
+
+    Returns:
+        A copy of ``x`` with the replacement values.
+    """
+    raise NotImplementedError("no `assign_sequence` method implemented for '" + type(x[0]).__name__ + "' objects")
+
+
+@assign_sequence.register
+def _assign_sequence_list(x: list, indices: Sequence[int], replacement: Any) -> list:
+    output = x.copy()
+    for i, j in enumerate(indices):
+        output[j] = replacement[i]
+    return output
+
+
+@assign_sequence.register
+def _assign_sequence_numpy(x: numpy.ndarray, indices: Sequence[int], replacement: Any) -> numpy.ndarray:
+    output = numpy.copy(x)
+    output[indices] = replacement
+    return output
+
+
+@assign_sequence.register
+def _assign_sequence_range(x: range, indices: Sequence[int], replacement: Any) -> Union[range, list]:
+    if isinstance(replacement, range) and isinstance(indices, range) and x[slice(indices.start, indices.stop, indices.step)] == replacement:
+        return x
+
+    output = list(x)
+    for i, j in enumerate(indices):
+        output[j] = replacement[i]
+    return output

tests/test_assign.py

@@ -0,0 +1,11 @@
+from biocutils import assign, assign_rows
+import numpy as np
+
+
+def test_assign_overall():
+    x = [1, 2, 3, 4, 5]
+    assert assign(x, [0, 2, 4], ["A", "B", "C"]) == ["A", 2, "B", 4, "C"]
+
+    y = np.random.rand(10, 20)
+    y2 = np.random.rand(5, 20)
+    assert (assign(y, range(5), y2) == assign_rows(y, range(5), y2)).all()

tests/test_assign_rows.py

@@ -0,0 +1,9 @@
+from biocutils import assign_rows
+import numpy as np
+
+
+def test_assign_numpy():
+    y = np.random.rand(10, 20)
+    y2 = np.random.rand(5, 20)
+    expected = np.concatenate([y2, y[5:10,:]])
+    assert (assign_rows(y, range(5), y2) == expected).all()

tests/test_assign_sequence.py

@@ -0,0 +1,20 @@
+from biocutils import assign_sequence
+import numpy as np
+
+
+def test_assign_list():
+    x = [1, 2, 3, 4, 5]
+    assert assign_sequence(x, [2,3,4], [0, 2, 4]) == [1, 2, 0, 2, 4]
+
+
+def test_assign_numpy():
+    y1 = np.random.rand(10)
+    y2 = np.random.rand(5)
+    expected = np.concatenate([y1[:5], y2])
+    assert (assign_sequence(y1, range(5, 10), y2) == expected).all()
+
+
+def test_assign_range():
+    x = range(10, 20)
+    assert assign_sequence(x, range(2, 7), ["A", "B", "C", "D", "E"]) == [10, 11, "A", "B", "C", "D", "E", 17, 18, 19]
+    assert assign_sequence(x, range(2, 7), range(12, 17)) == range(10, 20)