Repository: ul-mds/gecko

Issue:
Add mutator for appending or prepending data from another generator
This is useful for mimicking compound names or adding prefixes and suffixes that would otherwise not be there.

Discussion:

Solution patch:
diff --git a/docs/data-mutation.md b/docs/data-mutation.md
index 383375b..92aa5dd 100644
--- a/docs/data-mutation.md
+++ b/docs/data-mutation.md
@@ -425,6 +425,42 @@ print(datetime_mutator([srs]))
 # => ["2020-01-01", "2020-01-03", ..., "2020-01-25", "2020-01-31"]
 ```
 
+### Using generators
+
+`with_generator` can leverage Gecko's mutators to prepend, append or replace data.
+For instance, this can be used for emulating compound names for persons who have more than one given or last name.
+By default, this function adds a space when prepending or appending generated data, but this can be customized.
+
+```python
+import numpy as np
+import pandas as pd
+
+from gecko import mutator
+
+
+def generate_foobar_suffix(rand: np.random.Generator):
+    def _generate(count: int) -> list[pd.Series]:
+        return [pd.Series(rand.choice(("bar", "baz", "bat"), size=count))]
+
+    return _generate
+
+
+srs = pd.Series(["foo"] * 100)
+rng = np.random.default_rng(0x25504446)
+
+gen_prepend_mutator = mutator.with_generator(generate_foobar_suffix(rng), "prepend")
+print(gen_prepend_mutator([srs]))
+# => ["bat foo", "bar foo", ..., "baz foo", "baz foo"]
+
+gen_replace_mutator = mutator.with_generator(generate_foobar_suffix(rng), "replace")
+print(gen_replace_mutator([srs]))
+# => ["bar", "bar", ..., "baz", "bat"]
+
+gen_append_mutator = mutator.with_generator(generate_foobar_suffix(rng), "append", join_with="")
+print(gen_append_mutator([srs]))
+# => ["foobat", "foobat", ..., "foobat", "foobaz"]
+```
+
 ## Multiple mutators
 
 Using `mutate_data_frame`, you can apply multiple mutators on many columns at once.
diff --git a/gecko/mutator.py b/gecko/mutator.py
index 71191d6..e246983 100644
--- a/gecko/mutator.py
+++ b/gecko/mutator.py
@@ -21,6 +21,7 @@
     "with_lowercase",
     "with_uppercase",
     "with_datetime_offset",
+    "with_generator",
     "mutate_data_frame",
 ]
 
@@ -37,6 +38,7 @@
 from typing_extensions import ParamSpec, Concatenate
 
 from gecko.cldr import decode_iso_kb_pos, unescape_kb_char, get_neighbor_kb_pos_for
+from gecko.generator import Generator
 
 Mutator = Callable[[list[pd.Series]], list[pd.Series]]
 _EditOp = Literal["ins", "del", "sub", "trs"]
@@ -1292,6 +1294,76 @@ def _mutate(srs_lst: list[pd.Series]) -> list[pd.Series]:
     return _mutate
 
 
+def with_generator(
+    generator: Generator,
+    mode: Literal["prepend", "append", "replace"],
+    join_with: str = " ",
+) -> Mutator:
+    """
+    Mutate data from a series by appending, prepending or replacing it with data from another generator.
+    A character to join generated data with when appending or prepending can be provided.
+
+    Args:
+        generator: generator to source data from
+        mode: either append, prepend or replace
+        join_with: join character when appending or prepending
+
+    Returns:
+        function returning list of strings that have been appended, prepended or replaced with data from a generator
+    """
+
+    def _mutate(srs_lst: list[pd.Series]) -> list[pd.Series]:
+        # check that all series are of the same length
+        srs_lst_len_set = set([len(srs) for srs in srs_lst])
+
+        if len(srs_lst_len_set) != 1:
+            raise ValueError("series do not have the same length")
+
+        # use this length as a param for the generator later
+        srs_len = srs_lst_len_set.pop()
+
+        # check that the indices of all input series are aligned
+        if len(srs_lst) > 1:
+            indices_aligned = [
+                (srs_lst[0].index == srs_lst[i].index).all()
+                for i in range(1, len(srs_lst))
+            ]
+
+            if not all(indices_aligned):
+                raise ValueError("indices of input series are not aligned")
+
+        # call generator and align its index with the input series index. use ffill to
+        # avoid nas when reindexing.
+        srs_gen_lst = [
+            srs.reindex(srs_lst[i].index, method="ffill")
+            for i, srs in enumerate(generator(srs_len))
+        ]
+
+        # check that the generator returns as many series as provided to the mutator.
+        if len(srs_lst) != len(srs_gen_lst):
+            raise ValueError(
+                f"generator must generate as many series as provided to the mutator: "
+                f"got {len(srs_gen_lst)}, expected {len(srs_lst)}"
+            )
+
+        srs_lst_out = [srs.copy() for srs in srs_lst]
+
+        # perform the actual data mutation (this is where index alignment matters)
+        for i, srs_gen in enumerate(srs_gen_lst):
+            if mode == "replace":
+                srs_lst_out[i][:] = srs_gen
+            elif mode == "prepend":
+                srs_lst_out[i][:] = srs_gen + join_with + srs_lst_out[i][:]
+            elif mode == "append":
+                srs_lst_out[i][:] += join_with + srs_gen
+            else:
+                raise ValueError(f"invalid mode: `{mode}`")
+
+        return srs_lst_out
+
+    return _mutate
+
+
 def mutate_data_frame(
     df_in: pd.DataFrame,
     column_to_mutator_list: list[


Test patch:
diff --git a/tests/test_mutator.py b/tests/test_mutator.py
index 7ae2174..f16d55d 100644
--- a/tests/test_mutator.py
+++ b/tests/test_mutator.py
@@ -4,6 +4,7 @@
 import pandas as pd
 import pytest
 
+from gecko.generator import Generator
 from gecko.mutator import (
     with_insert,
     with_delete,
@@ -23,6 +24,7 @@
     with_lowercase,
     with_uppercase,
     with_datetime_offset,
+    with_generator,
 )
 from tests.helpers import get_asset_path
 
@@ -397,6 +399,90 @@ def test_with_datetime_offset_raise_nonpositive_delta(rng):
     assert str(e.value) == "delta must be positive, is 0"
 
 
+def _generate_static(value: str) -> Generator:
+    def _generate(count: int) -> list[pd.Series]:
+        return [pd.Series([value] * count)]
+
+    return _generate
+
+
+def test_with_generator_replace():
+    srs = pd.Series(["foo"] * 100)
+    mutate_generator = with_generator(_generate_static("bar"), "replace")
+    (srs_mutated,) = mutate_generator([srs])
+
+    assert (srs_mutated == "bar").all()
+
+
+def test_with_generator_prepend():
+    srs = pd.Series(["foo"] * 100)
+    mutate_generator = with_generator(_generate_static("bar"), "prepend")
+    (srs_mutated,) = mutate_generator([srs])
+
+    assert (srs_mutated == "bar foo").all()
+
+
+def test_with_generator_append():
+    srs = pd.Series(["foo"] * 100)
+    mutate_generator = with_generator(_generate_static("bar"), "append")
+    (srs_mutated,) = mutate_generator([srs])
+
+    assert (srs_mutated == "foo bar").all()
+
+
+def test_with_generator_join_character():
+    srs = pd.Series(["foo"] * 100)
+    mutate_generator = with_generator(_generate_static("bar"), "append", join_with="-")
+    (srs_mutated,) = mutate_generator([srs])
+
+    assert (srs_mutated == "foo-bar").all()
+
+
+def test_with_generator_slice():
+    srs = pd.Series(["foo"] * 100)
+    srs_sub = srs.iloc[range(0, len(srs), 3)]
+    mutate_generator = with_generator(_generate_static("bar"), "append")
+    (srs_mutated,) = mutate_generator([srs_sub])
+
+    assert (srs_mutated == "foo bar").all()
+
+
+def test_with_generator_raise_series_different_length():
+    srs_1, srs_2 = pd.Series(["foo"] * 100), pd.Series(["bar"] * 200)
+    mutate_generator = with_generator(_generate_static("bar"), "append")
+
+    with pytest.raises(ValueError) as e:
+        mutate_generator([srs_1, srs_2])
+
+    assert str(e.value) == "series do not have the same length"
+
+
+def test_with_generator_raise_generator_incompatible():
+    srs_1, srs_2 = pd.Series(["foo"] * 100), pd.Series(["bar"] * 100)
+    mutate_generator = with_generator(_generate_static("bar"), "append")
+
+    with pytest.raises(ValueError) as e:
+        mutate_generator([srs_1, srs_2])
+
+    assert (
+        str(e.value)
+        == "generator must generate as many series as provided to the mutator: got 1, expected 2"
+    )
+
+
+def test_with_generator_raise_unaligned_indices():
+    srs_1, srs_2 = (
+        pd.Series(["foo"] * 3, index=[0, 1, 2]),
+        pd.Series(["bar"] * 3, index=[3, 4, 5]),
+    )
+    mutate_generator = with_generator(_generate_static("bar"), "append")
+
+    with pytest.raises(ValueError) as e:
+        mutate_generator([srs_1, srs_2])
+
+    assert str(e.value) == "indices of input series are not aligned"
+
+
 def test_mutate_data_frame_single(rng):
     df = pd.DataFrame({"foo": list(string.ascii_letters)})
     df_mut = mutate_data_frame(
@@ -591,6 +677,9 @@ def test_mutate_data_frame_multicolumn_noop():
         (2, with_permute()),
         (1, with_lowercase()),
         (1, with_uppercase()),
+        (1, with_generator(_generate_static("foo"), "prepend")),
+        (1, with_generator(_generate_static("foo"), "append")),
+        (1, with_generator(_generate_static("foo"), "replace")),
     ],
 )
 def test_mutator_no_modify(num_srs: int, func: Mutator, rng):