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.

Hints from issue comments:

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):