install-stories command allows story file to specify its configuration

.gitignore

@@ -151,7 +151,5 @@ docs/temp/*
 /src-stats.yaml
 /config.yaml
 *.yaml.gz
-
-*_stories.py
-
 /examples/opiates/*
+./*_stories.py

datafaker/install.py

@@ -0,0 +1,244 @@
+"""Functions to install Python file references in ``config.yaml``."""
+from collections.abc import Mapping, MutableMapping, Sequence
+from inspect import Parameter, signature
+from pathlib import Path
+from typing import Any
+
+from datafaker.utils import import_file, logger
+
+
+def _make_where_from_annotation(
+    query_def: Mapping[str, Any],
+    fn_name: str,
+    param_name: str,
+) -> str:
+    """Make a where clause from ``query`` value from the annotation."""
+    if "where" not in query_def:
+        return ""
+    w = query_def["where"]
+    if isinstance(w, str):
+        return f" WHERE {w}"
+    if isinstance(w, Sequence):
+        return " WHERE " + " AND ".join(f'"({clause})"' for clause in w)
+    logger.warning(
+        '"where" in the query annotation of parameter "%s" of function "%s"'
+        " needs to be a string or a list of strings",
+        param_name,
+        fn_name,
+    )
+    return ""
+
+
+def _make_vars_from_annotation(
+    query_def: Mapping[str, Any],
+    fn_name: str,
+    param_name: str,
+) -> Mapping[str, Any]:
+    """Make a variables dict from ``query`` value from the annotation."""
+    if "vars" not in query_def:
+        return {}
+    vars_def = query_def["vars"]
+    if isinstance(vars_def, Mapping):
+        return vars_def
+    if isinstance(vars_def, Sequence):
+        return {v: v for v in query_def["vars"]}
+    logger.warning(
+        '"vars" in the query annotation of parameter "%s" of function "%s"'
+        " needs to be a list of strings or a dict of strings to strings",
+        param_name,
+        fn_name,
+    )
+    return {}
+
+
+def _add_count_vars_from_annotation(
+    group_vars_out: MutableMapping[str, Any],
+    query_def: Mapping[str, Any],
+    fn_name: str,
+    param_name: str,
+) -> None:
+    """Add ``GROUP BY`` clauses from ``count_vars``."""
+    if "count_vars" not in query_def:
+        return
+    cntv = query_def["count_vars"]
+    if isinstance(cntv, Mapping):
+        group_vars_out.update({k: f"COUNT({v})" for k, v in cntv})
+        return
+    logger.warning(
+        '"count_vars" needs to be a dict in the annotation for parameter %s of function %s',
+        param_name,
+        fn_name,
+    )
+
+
+def _add_ms_vars_from_annotation(
+    group_vars_out: MutableMapping[str, Any],
+    query_def: Mapping[str, Any],
+    fn_name: str,
+    param_name: str,
+) -> None:
+    """Add ``GROUP BY`` clauses from ``ms_vars``."""
+    if "ms_vars" not in query_def:
+        return
+    msv = query_def["ms_vars"]
+    if not isinstance(msv, Mapping):
+        logger.warning(
+            '"ms_vars" needs to be a dict in the annotation for parameter %s of function %s',
+            param_name,
+            fn_name,
+        )
+        return
+    for k, v in msv.items():
+        group_vars_out[k + "_count"] = f"COUNT({v})"
+        group_vars_out[k + "_mean"] = f"AVG({v})"
+        group_vars_out[k + "_stddev"] = f"STDDEV({v})"
+
+
+def make_query_from_annotation(
+    annotation_data: Any,
+    fn_name: str,
+    param_name: str,
+) -> str | None:
+    """
+    Make new configuration items describing a query.
+
+    The query's result will be passed as this parameter to this function.
+
+    The annotation must be a dict with the following keys:
+
+    ``comment``: A string describing the query in natural language.
+
+    ``query``: Either a string containing the SQL query required, or
+    a dict containing the following keys:
+
+    * ``table``: The table to query. Could be "tablename AS alias" if you like.
+    * ``vars`` (optional): Either a list of columns to extract from the table(s),
+      or a dict of keys (the names of the keys in the dict to be passed to the
+      annotated function) to values (the names of the columns to be extracted).
+      At least one of ``vars``, ``ms_vars``, ``count_vars`` must be present.
+    * ``where`` (optional): A SQL expression to filter the results.
+    * ``count_vars`` (optional): A dict of keys to be passed to the function
+      to values that are the names of the columns to be counted (could be
+      ``*``; if the name of a column the result will be the number of non-null
+      entries in that column). The query will be grouped by ``vars``.
+    * ``ms_vars`` (optional): A dict of value names to columns to be analysed.
+      The keys to be passed to the function will be name + ``_count`` for the
+      number of non-null values in that column, name + ``_mean`` for the
+      average value in that column and name + ``_stddev`` for the standard
+      deviation of values in that column.
+
+    :param annotation_data: The ``Annotation`` attached to the parameter.
+    :param fn_name: The name of the function that the parameter is of.
+    :param param_name: The name of the parameter with the annotation.
+    :return: A mapping of new configuration items to add to the configuration,
+    if the annotation had a well-defined query and comment value; otherwise
+    an empty dict.
+    """
+    if not isinstance(annotation_data, Sequence):
+        return None
+    ann = annotation_data[0]
+    if not isinstance(ann, Mapping) or "query" not in ann:
+        return None
+    if isinstance(ann["query"], str):
+        return ann["query"]
+    query_def = ann["query"]
+    if "table" not in query_def:
+        logger.warning(
+            '"table" needs to be a key in the annotation for'
+            ' the "query" value of parameter "%s" of function "%s"',
+            param_name,
+            fn_name,
+        )
+        return None
+    table = query_def["table"]
+    nongroup_vars = _make_vars_from_annotation(query_def, fn_name, param_name)
+    where = _make_where_from_annotation(query_def, fn_name, param_name)
+    group_vars: dict[str, Any] = {}
+    _add_count_vars_from_annotation(group_vars, query_def, fn_name, param_name)
+    _add_ms_vars_from_annotation(group_vars, query_def, fn_name, param_name)
+    if group_vars and nongroup_vars:
+        group_by = " GROUP BY " + ", ".join(f'"{v}"' for v in nongroup_vars)
+    else:
+        group_by = ""
+    vars_exprs = ", ".join(
+        f'{v} AS "{k}"' for k, v in {**nongroup_vars, **group_vars}.items()
+    )
+    return f"SELECT {vars_exprs} FROM {table}{group_by}{where}"
+
+
+def _add_kwarg(
+    kwargs_out: dict[str, Any], fn_name: str, param: Parameter
+) -> list[dict[str, Any]]:
+    """
+    Add a kwargs configuration and return a ``src_stats`` query item.
+
+    :param kwargs_out: The story generator's ``kwargs`` value to be updated.
+    :param fn_name: The name of the story generator function.
+    :param param: The parameter to specify.
+    :return: A list of configuration items to add to the ``src_stats`` config, for
+    all the queries this parameter requires.
+    """
+    if param.annotation is Parameter.empty:
+        return []
+    meta = param.annotation.__metadata__
+    query = make_query_from_annotation(
+        param.annotation.__metadata__, fn_name, param.name
+    )
+    if query is None:
+        return []
+    stat_name = f"story_auto__{fn_name}__{param.name}"
+    if "comments" in meta[0]:
+        comments = [meta[0]["comment"]]
+    else:
+        comments = []
+    ssc = {
+        "name": stat_name,
+        "query": query,
+        "comments": comments,
+    }
+    kwargs_out[param.name] = f'SRC_STATS["{stat_name}"]["results"]'
+    return [ssc]
+
+
+def install_stories_from(config: MutableMapping[str, Any], story_file: Path) -> bool:
+    """
+    Configure datafaker with the stories in a Python file.
+
+    :param config: The contents of the configuration file, to be mutated.
+    :param story_file: Path to the Python file containing the story generators.
+    :return: True if the config was updated correctly, False if it was untouched
+    because problems were encountered.
+    """
+    story_generators: list[Mapping[str, Any]] = []
+    src_stats = [
+        s
+        for s in config.get("src_stats", [])
+        if isinstance(s, Mapping)
+        and "name" in s
+        and not s["name"].startswith("story_auto__")
+    ]
+    story_module_name = story_file.stem
+    story_module = import_file(story_file, story_module_name)
+    for attr_name in dir(story_module):
+        attr = getattr(story_module, attr_name)
+        if (
+            hasattr(attr, "__module__")
+            and attr.__module__ == story_module_name
+            and not attr_name.startswith("_")
+            and callable(attr)
+        ):
+            kwargs: dict[str, None] = {}
+            sig = signature(attr)
+            for param in sig.parameters.values():
+                src_stats += _add_kwarg(kwargs, attr, param)
+            story_generators.append(
+                {
+                    "name": f"{story_module_name}.{attr_name}",
+                    "num_stories_per_pass": 1,
+                    "kwargs": kwargs,
+                }
+            )
+    config["story_generators_module"] = story_module_name
+    config["story_generators"] = story_generators
+    config["src-stats"] = src_stats
+    return True

datafaker/main.py

@@ -24,13 +24,18 @@
     TableWriter,
     get_parquet_table_writer,
 )
+from datafaker.install import install_stories_from
 from datafaker.interactive import (
     update_config_generators,
     update_config_tables,
     update_missingness,
 )
 from datafaker.interactive.base import DbCmd
-from datafaker.make import make_src_stats, make_tables_file, make_vocabulary_tables
+from datafaker.make import (
+    make_src_stats,
+    make_tables_file,
+    make_vocabulary_tables,
+)
 from datafaker.remove import remove_db_data, remove_db_tables, remove_db_vocab
 from datafaker.settings import (
     SettingsError,
@@ -265,7 +270,7 @@ def create_tables(
         $ datafaker create-tables
     """
     logger.debug("Creating tables.")
-    config = read_config_file(config_file) if config_file is not None else {}
+    config = read_config_file(config_file)
     orm_metadata = load_metadata_for_output(orm_file, config)
     create_db_tables(orm_metadata)
     logger.debug("Tables created.")
@@ -710,7 +715,7 @@ def remove_data(
     """Truncate non-vocabulary tables in the destination schema."""
     if yes:
         logger.debug("Truncating non-vocabulary tables.")
-        config = read_config_file(config_file) if config_file is not None else {}
+        config = read_config_file(config_file)
         metadata = load_metadata_for_output(orm_file, config)
         remove_db_data(metadata, config)
         logger.debug("Non-vocabulary tables truncated.")
@@ -737,7 +742,7 @@ def remove_vocab(
     """Truncate vocabulary tables in the destination schema."""
     if yes:
         logger.debug("Truncating vocabulary tables.")
-        config = read_config_file(config_file) if config_file is not None else {}
+        config = read_config_file(config_file)
         meta_dict = load_metadata_config(orm_file, config)
         orm_metadata = dict_to_metadata(meta_dict, config)
         remove_db_vocab(orm_metadata, meta_dict, config)
@@ -812,7 +817,7 @@ def list_tables(
     tables: TableType = Option(TableType.GENERATED, help="Which tables to list"),
 ) -> None:
     """List the names of tables described in the metadata file."""
-    config = read_config_file(config_file) if config_file is not None else {}
+    config = read_config_file(config_file)
     orm_metadata = load_metadata(orm_file, config)
     all_table_names = set(orm_metadata.tables.keys())
     vocab_table_names = {
@@ -830,6 +835,26 @@ def list_tables(
         print(name)
 
 
+@app.command()
+def install_stories(
+    config_file: Path = Option(CONFIG_FILENAME, help="The configuration file"),
+    story_file: Path = Argument(help="The Python file containing stories"),
+) -> None:
+    """Add the story file's name and any contained query to the configuration file."""
+    config_file_path = Path(config_file)
+    config = {}
+    if config_file_path.exists():
+        config = yaml.load(
+            config_file_path.read_text(encoding="UTF-8"), Loader=yaml.SafeLoader
+        )
+    if not install_stories_from(config, story_file):
+        logger.debug("Cancelled")
+        sys.exit(1)
+    content = yaml.dump(config)
+    config_file_path.write_text(content, encoding="utf-8")
+    logger.debug("Stories configured in %s.", config_file)
+
+
 @app.command()
 def version() -> None:
     """Display version information."""

datafaker/utils.py

@@ -67,7 +67,7 @@ def read_config_file(path: Path) -> dict:
     return config
 
 
-def import_file(file_path: str) -> ModuleType:
+def import_file(file_path: str | Path, module_name: str = "df") -> ModuleType:
     """Import a file.
 
     This utility function returns file_path imported as a module.
@@ -78,7 +78,7 @@ def import_file(file_path: str) -> ModuleType:
     Returns:
         ModuleType
     """
-    spec = importlib.util.spec_from_file_location("df", file_path)
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
     if spec is None or spec.loader is None:
         raise SettingsError(f"No loadable module '{file_path}'")
     module = importlib.util.module_from_spec(spec)

examples/mimic_omop/README.md

@@ -8,10 +8,6 @@
 
 `poetry run datafaker create-tables --orm-file ./examples/mimic_omop/orm.yaml --config-file ./examples/mimic_omop/config.yaml`
 
-1. Create generator table
-
-`poetry run datafaker create-generators --orm-file ./examples/mimic_omop/orm.yaml --config-file ./examples/mimic_omop/config.yaml --df-file ./examples/mimic_omop/df.py`
-
 1. Create data
 
 `poetry run datafaker create-data --orm-file ./examples/mimic_omop/orm.yaml --config-file ./examples/mimic_omop/config.yaml --df-file .\examples\mimic_omop\df.py`

tests/examples/annotated_stories.py

@@ -0,0 +1,28 @@
+"""Story generators which describe their own queries and can therefore be installed."""
+from collections.abc import Iterable
+from typing import Annotated, Any
+
+def string_story_one_sd(
+    stats: Annotated[dict, {
+        "query": {
+            "ms_vars": {"freq": "frequency"},
+            "table": "string",
+        },
+        "comment": "Frequency mean and standard deviation",
+    }],
+) -> Iterable[tuple[str, dict[str, Any]]]:
+    man = yield("manufacturer", {"name": "one"})
+    mod = yield ("model", {
+        "name": "one_sd",
+        "manufacturer_id": man["id"]
+    })
+    yield("string", {
+        "model_id": mod["id"],
+        "position": 0,
+        "frequency": stats[0]["freq_mean"] - stats[0]["freq_stddev"],
+    })
+    yield("string", {
+        "model_id": mod["id"],
+        "position": stats[0]["freq_count"],
+        "frequency": stats[0]["freq_mean"] + stats[0]["freq_stddev"],
+    })

tests/examples/install_config.yaml

@@ -0,0 +1,3 @@
+tables:
+  string:
+    num_rows_per_pass: 0