Repository: ul-mds/gecko

Issue:
Allow data frames to be passed to any function reading CSV files
This affects the following functions:

- [ ] `generator.from_frequency_table`
- [ ] `generator.from_multicolumn_frequency_table`
- [ ] `mutator.with_phonetic_replacement_table`
- [ ] `mutator.with_replacement_table`
- [ ] `mutator.with_categorical_values`
- [ ] `mutator.with_regex_replacement_table`

Discussion:

Solution patch:
diff --git a/gecko/generator.py b/gecko/generator.py
index c083ef0..08f9b5b 100644
--- a/gecko/generator.py
+++ b/gecko/generator.py
@@ -110,7 +110,7 @@ def _generate(count: int) -> list[pd.Series]:
 
 
 def from_frequency_table(
-    csv_file_path: _t.Union[str, PathLike[str]],
+    data_source: _t.Union[str, PathLike[str], pd.DataFrame],
     value_column: _t.Union[str, int] = 0,
     freq_column: _t.Union[str, int] = 1,
     encoding: str = "utf-8",
@@ -126,7 +126,7 @@ def from_frequency_table(
     has a header row.
 
     Args:
-        csv_file_path: path to CSV file
+        data_source: path to CSV file or data frame to use as frequency table
         value_column: name or index of the value column
         freq_column: name or index of the frequency column
         encoding: character encoding of the CSV file
@@ -143,23 +143,26 @@ def from_frequency_table(
         raise ValueError("value and frequency columns must be of the same type")
 
     # skip check for value_column bc they are both of the same type already
-    if not isinstance(freq_column, str) and not isinstance(freq_column, int):
+    if not isinstance(freq_column, (str, int)):
         raise ValueError(
             "value and frequency columns must be either a string or an integer"
         )
 
-    header = isinstance(freq_column, str)
-
-    # read csv file
-    df = pd.read_csv(
-        csv_file_path,
-        header=0 if header else None,  # header row index (`None` if not present)
-        usecols=[value_column, freq_column],
-        dtype={freq_column: "int", value_column: "str"},
-        keep_default_na=False,
-        sep=delimiter,
-        encoding=encoding,
-    )
+    if isinstance(data_source, pd.DataFrame):
+        df = data_source
+    else:
+        header = isinstance(freq_column, str)
+
+        # read csv file
+        df = pd.read_csv(
+            data_source,
+            header=0 if header else None,  # header row index (`None` if not present)
+            usecols=[value_column, freq_column],
+            dtype={freq_column: "int", value_column: "str"},
+            keep_default_na=False,
+            sep=delimiter,
+            encoding=encoding,
+        )
 
     # convert absolute to relative frequencies
     srs_value = df[value_column]
@@ -172,7 +175,7 @@ def _generate(count: int) -> list[pd.Series]:
 
 
 def from_multicolumn_frequency_table(
-    csv_file_path: _t.Union[str, PathLike[str]],
+    data_source: _t.Union[str, PathLike[str], pd.DataFrame],
     value_columns: _t.Union[int, str, list[int], list[str]] = 0,
     freq_column: _t.Union[int, str] = 1,
     encoding: str = "utf-8",
@@ -188,7 +191,7 @@ def from_multicolumn_frequency_table(
     has a header row.
 
     Args:
-        csv_file_path: path to CSV file
+        data_source: path to CSV file or data frame to use as frequency table
         value_columns: names or indices of the value columns
         freq_column: name or index of the frequency column
         encoding: character encoding of the CSV file
@@ -227,20 +230,23 @@ def from_multicolumn_frequency_table(
                 "value and frequency column must be either a string or an integer"
             )
 
-    header = isinstance(freq_column, str)
-
-    df = pd.read_csv(
-        csv_file_path,
-        header=0 if header else None,
-        usecols=value_columns + [freq_column],
-        dtype={
-            freq_column: "int",
-            **{value_column: "str" for value_column in value_columns},
-        },
-        keep_default_na=False,
-        sep=delimiter,
-        encoding=encoding,
-    )
+    if isinstance(data_source, pd.DataFrame):
+        df = data_source
+    else:
+        header = isinstance(freq_column, str)
+
+        df = pd.read_csv(
+            data_source,
+            header=0 if header else None,
+            usecols=value_columns + [freq_column],
+            dtype={
+                freq_column: "int",
+                **{value_column: "str" for value_column in value_columns},
+            },
+            keep_default_na=False,
+            sep=delimiter,
+            encoding=encoding,
+        )
 
     # sum of absolute frequencies
     freq_total = df[freq_column].sum()
diff --git a/gecko/mutator.py b/gecko/mutator.py
index 0628738..e7326fb 100644
--- a/gecko/mutator.py
+++ b/gecko/mutator.py
@@ -262,7 +262,7 @@ def _mutate(srs_lst: list[pd.Series]) -> list[pd.Series]:
 
 
 def with_phonetic_replacement_table(
-    csv_file_path: _t.Union[PathLike, str],
+    data_source: _t.Union[PathLike, str, pd.DataFrame],
     source_column: _t.Union[int, str] = 0,
     target_column: _t.Union[int, str] = 1,
     flags_column: _t.Union[int, str] = 2,
@@ -283,7 +283,7 @@ def with_phonetic_replacement_table(
     has a header row.
 
     Args:
-        csv_file_path: path to CSV file containing phonetic replacement rules
+        data_source: path to CSV file or data frame containing phonetic replacement rules
         source_column: name or index of source column
         target_column: name or index of target column
         flags_column: name or index of flag column
@@ -326,18 +326,21 @@ def __new_unknown_flag_error(flag: str):
             "source, target and flags columns must be either a string or an integer"
         )
 
-    header = isinstance(flags_column, str)
-
-    # read csv file
-    df = pd.read_csv(
-        csv_file_path,
-        header=0 if header else None,
-        dtype=str,
-        usecols=[source_column, target_column, flags_column],
-        keep_default_na=False,
-        sep=delimiter,
-        encoding=encoding,
-    )
+    if isinstance(data_source, pd.DataFrame):
+        df = data_source
+    else:
+        header = isinstance(flags_column, str)
+
+        # read csv file
+        df = pd.read_csv(
+            data_source,
+            header=0 if header else None,
+            dtype=str,
+            usecols=[source_column, target_column, flags_column],
+            keep_default_na=False,
+            sep=delimiter,
+            encoding=encoding,
+        )
 
     # parse replacement rules
     phonetic_replacement_rules: list[_PhoneticReplacementRule] = []
@@ -473,7 +476,7 @@ def _mutate(srs_lst: list[pd.Series]) -> list[pd.Series]:
 
 
 def with_replacement_table(
-    csv_file_path: _t.Union[PathLike, str],
+    data_source: _t.Union[PathLike, str, pd.DataFrame],
     source_column: _t.Union[str, int] = 0,
     target_column: _t.Union[str, int] = 1,
     inline: bool = False,
@@ -494,7 +497,7 @@ def with_replacement_table(
     has a header row.
 
     Args:
-        csv_file_path: path to CSV file
+        data_source: path to CSV file or data frame containing replacement table
         source_column: name or index of the source column
         target_column: name or index of the target column
         inline: whether to perform replacements inline
@@ -519,17 +522,20 @@ def with_replacement_table(
             "source and target columns must be either a string or an integer"
         )
 
-    header = isinstance(target_column, str)
+    if isinstance(data_source, pd.DataFrame):
+        df = data_source
+    else:
+        header = isinstance(target_column, str)
 
-    df = pd.read_csv(
-        csv_file_path,
-        header=0 if header else None,
-        dtype=str,
-        usecols=[source_column, target_column],
-        keep_default_na=False,
-        sep=delimiter,
-        encoding=encoding,
-    )
+        df = pd.read_csv(
+            data_source,
+            header=0 if header else None,
+            dtype=str,
+            usecols=[source_column, target_column],
+            keep_default_na=False,
+            sep=delimiter,
+            encoding=encoding,
+        )
 
     if reverse:
         # flip columns and concat
@@ -1038,7 +1044,7 @@ def _mutate(srs_lst: list[pd.Series]) -> list[pd.Series]:
 
 
 def with_categorical_values(
-    csv_file_path: _t.Union[PathLike, str],
+    data_source: _t.Union[PathLike, str, pd.DataFrame],
     value_column: _t.Union[str, int] = 0,
     encoding: str = "utf-8",
     delimiter: str = ",",
@@ -1052,7 +1058,7 @@ def with_categorical_values(
     has a header row.
 
     Args:
-        csv_file_path: path to CSV file
+        data_source: path to CSV file or data frame containing values
         value_column: name or index of value column
         encoding: character encoding of the CSV file
         delimiter: column delimiter of the CSV file
@@ -1067,18 +1073,21 @@ def with_categorical_values(
     if not isinstance(value_column, str) and not isinstance(value_column, int):
         raise ValueError("value column must be either a string or an integer")
 
-    header = isinstance(value_column, str)
-
-    # read csv file
-    df = pd.read_csv(
-        csv_file_path,
-        header=0 if header else None,
-        dtype=str,
-        usecols=[value_column],
-        keep_default_na=False,
-        sep=delimiter,
-        encoding=encoding,
-    )
+    if isinstance(data_source, pd.DataFrame):
+        df = data_source
+    else:
+        header = isinstance(value_column, str)
+
+        # read csv file
+        df = pd.read_csv(
+            data_source,
+            header=0 if header else None,
+            dtype=str,
+            usecols=[value_column],
+            keep_default_na=False,
+            sep=delimiter,
+            encoding=encoding,
+        )
 
     # fetch unique values
     unique_values = pd.Series(df[value_column].unique())
@@ -1421,7 +1430,7 @@ def _parse_regex_flags(regex_flags_val: str) -> int:
 
 
 def with_regex_replacement_table(
-    csv_file_path: _t.Union[PathLike, str],
+    data_source: _t.Union[PathLike, str, pd.DataFrame],
     pattern_column: str = "pattern",
     flags_column: _t.Optional[str] = None,
     encoding: str = "utf-8",
@@ -1436,7 +1445,7 @@ def with_regex_replacement_table(
     When using named capture groups, the columns must be named after the capture groups they are supposed to substitute.
 
     Args:
-        csv_file_path: path to CSV file
+        data_source: path to CSV file or data frame
         pattern_column: name of regex pattern column
         flags_column: name of regex flag column
         encoding: character encoding of the CSV file
@@ -1449,16 +1458,19 @@ def with_regex_replacement_table(
     if rng is None:
         rng = np.random.default_rng()
 
-    df = pd.read_csv(
-        csv_file_path,
-        encoding=encoding,
-        keep_default_na=False,
-        sep=delimiter,
-        dtype=str,
-    )
+    if isinstance(data_source, pd.DataFrame):
+        df = data_source
+    else:
+        df = pd.read_csv(
+            data_source,
+            encoding=encoding,
+            keep_default_na=False,
+            sep=delimiter,
+            dtype=str,
+        )
 
     if pattern_column not in df.columns:
-        raise ValueError(f"CSV file at `{csv_file_path}` doesn't have a pattern column")
+        raise ValueError(f"CSV file at `{data_source}` doesn't have a pattern column")
 
     regexes: list[re.Pattern] = []
     regex_repl_fns: list[_t.Callable[[re.Match], str]] = []


Test patch:
diff --git a/tests/test_generator.py b/tests/test_generator.py
index 3a01876..35a9f83 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -244,6 +244,20 @@ def test_from_frequency_table_nan(tmp_path, rng):
     assert pd.notna(srs).all()
 
 
+def test_from_frequency_table_df(rng):
+    df = pd.DataFrame.from_dict({"freq": [10, 5], "value": ["foo", "bar"]})
+
+    gen_freq_table = generator.from_frequency_table(
+        df, value_column="value", freq_column="freq", rng=rng
+    )
+
+    (srs,) = gen_freq_table(100)
+    assert pd.notna(srs).all()
+
+    value_counts = srs.value_counts()
+    assert value_counts["bar"] < value_counts["foo"]
+
+
 def test_from_multicolumn_frequency_table_nan(tmp_path, rng):
     freq_file_path = tmp_path / "freq.csv"
     freq_file_path.write_text('value1,value2,freq\n"","bar",1\n"foo","baz",1\n')
@@ -259,3 +273,30 @@ def test_from_multicolumn_frequency_table_nan(tmp_path, rng):
 
     assert pd.notna(srs1).all()
     assert pd.notna(srs2).all()
+
+
+def test_from_multicolumn_frequency_table_df(rng):
+    df = pd.DataFrame.from_dict(
+        {
+            "freq": [10, 5, 20, 10],
+            "value1": ["foo", "foo", "bar", "bar"],
+            "value2": ["baz", "bat", "baz", "bat"],
+        }
+    )
+
+    gen_freq_table = generator.from_multicolumn_frequency_table(
+        df,
+        value_columns=["value1", "value2"],
+        freq_column="freq",
+        rng=rng,
+    )
+
+    (srs_1, srs_2) = gen_freq_table(100)
+    assert pd.notna(srs_1).all()
+    assert pd.notna(srs_2).all()
+
+    counts_1 = srs_1.value_counts()
+    assert counts_1["bar"] > counts_1["foo"]
+
+    counts_2 = srs_2.value_counts()
+    assert counts_2["baz"] > counts_2["bat"]
diff --git a/tests/test_mutator.py b/tests/test_mutator.py
index 3a59bdf..54516a3 100644
--- a/tests/test_mutator.py
+++ b/tests/test_mutator.py
@@ -189,6 +189,30 @@ def _generate_gender_list():
     assert ~(srs == srs_mutated).all()
 
 
+def test_with_categorical_values_df(rng):
+    def _generate_gender_list():
+        nonlocal rng
+        return rng.choice(["m", "f", "d", "x"], size=1000)
+
+    df_categorical = pd.DataFrame.from_dict(
+        {"id": list(range(16)), "gender": list("mfdx") * 4}, dtype=str
+    )
+
+    mutate_categorical = with_categorical_values(
+        df_categorical,
+        value_column="gender",
+        rng=rng,
+    )
+
+    srs = pd.Series(_generate_gender_list())
+    (srs_mutated,) = mutate_categorical([srs])
+
+    # same length
+    assert len(srs) == len(srs_mutated)
+    # different items
+    assert ~(srs == srs_mutated).all()
+
+
 def test_with_edit(rng):
     def _new_string():
         nonlocal rng
@@ -235,6 +259,40 @@ def test_with_phonetic_replacement_table(rng):
     assert (srs_mutated_actual == srs_mutated_expected).all()
 
 
+def test_with_phonetic_replacement_table_df(rng):
+    df_phonetic_in_out = pd.read_csv(get_asset_path("phonetic-test.csv"))
+    srs_original = df_phonetic_in_out["original"]
+    srs_mutated_expected = df_phonetic_in_out["corrupt"]
+
+    df_phonetic_rules = pd.DataFrame.from_records(
+        [
+            ("sch", "sh", ""),
+            ("ie", "i", ""),
+            ("ß", "ss", ""),
+            ("fi", "vi", ""),
+            ("ee", "eh", ""),
+            ("dt", "tt", "$"),
+            ("ah", "a", ""),
+            ("ph", "f", ""),
+            ("co", "ko", ""),
+            ("fe", "ve", "^"),
+            ("ei", "ai", ""),
+            ("uhr", "ur", ""),
+            ("ur", "uhr", ""),
+            ("ied", "id", ""),
+            ("ca", "ka", ""),
+            ("ci", "zi", ""),
+            ("ka", "ca", ""),
+            ("zi", "ci", ""),
+        ]
+    )
+
+    mutate_phonetic = with_phonetic_replacement_table(df_phonetic_rules, rng=rng)
+    (srs_mutated_actual,) = mutate_phonetic([srs_original])
+
+    assert (srs_mutated_actual == srs_mutated_expected).all()
+
+
 def test_with_cldr_keymap_file(rng):
     srs = pd.Series(["d", "e"])
     mutate_cldr = with_cldr_keymap_file(get_asset_path("de-t-k0-windows.xml"), rng=rng)
@@ -293,6 +351,39 @@ def test_with_replacement_table(rng):
     assert (srs[msk_juan] == srs_mutated[msk_juan]).all()
 
 
+def test_with_replacement_table_df(rng):
+    srs = pd.Series(["Jan", "Jann", "Juan"] * 10)
+    df_replacement = pd.DataFrame.from_dict(
+        {
+            "source": ["Jan"] * 4 + ["Jann"] * 4,
+            "target": [
+                "Jann",
+                "Jean",
+                "John",
+                "Juan",
+                "Jean",
+                "Johann",
+                "John",
+                "Juan",
+            ],
+        }
+    )
+
+    mutate_replacement = with_replacement_table(
+        df_replacement,
+        source_column="source",
+        target_column="target",
+        rng=rng,
+    )
+    (srs_mutated,) = mutate_replacement([srs])
+
+    msk_juan = srs == "Juan"
+
+    assert len(srs) == len(srs_mutated)
+    assert (srs[~msk_juan] != srs_mutated[~msk_juan]).all()
+    assert (srs[msk_juan] == srs_mutated[msk_juan]).all()
+
+
 def test_with_replacement_table_reverse(rng):
     srs = pd.Series(["Jan", "Jann", "Juan"] * 10)
     mutate_replacement = with_replacement_table(
@@ -987,6 +1078,23 @@ def test_with_replacement_table_nan(tmp_path, rng):
     assert (srs_mut == ["foobar", "foobaz", "foobat"]).all()
 
 
+def test_with_regex_replacement_table(rng):
+    df_regex = pd.DataFrame.from_dict(
+        {"pattern": [".*(?P<dash>-).*"], "dash": [""]}, dtype=str
+    )
+
+    mut_replace = with_regex_replacement_table(
+        df_regex,
+        pattern_column="pattern",
+        rng=rng,
+    )
+
+    srs = pd.Series(["foo-bar", "foo-baz", "foo-bat"])
+    (srs_mut,) = mut_replace([srs])
+
+    assert (srs_mut == ["foobar", "foobaz", "foobat"]).all()
+
+
 def test_with_regex_replacement_table_nan(tmp_path, rng):
     replacement_table_file_path = tmp_path / "replacement.csv"
     replacement_table_file_path.write_text('pattern,dash\n".*(?P<dash>-).*",""\n')