From dfaa5eddcb76dc83a52e91e4fa9e42500f1dc615 Mon Sep 17 00:00:00 2001 From: AlexandruAchilleF Date: Sun, 11 Oct 2026 17:11:30 +0300 Subject: [PATCH] feat(cli): add create-tests command for autogenerating unit test stubs (#74) --- hamilton/cli/__main__.py | 29 +++++++++++ hamilton/cli/commands.py | 5 ++ hamilton/cli/logic.py | 105 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 139 insertions(+) diff --git a/hamilton/cli/__main__.py b/hamilton/cli/__main__.py index 58d1174fd..054f4fc28 100644 --- a/hamilton/cli/__main__.py +++ b/hamilton/cli/__main__.py @@ -323,6 +323,35 @@ def ui( config_file=config_file, ) +@cli.command() +def create_tests( + ctx: typer.Context, + module_path: Path = typer.Argument( + ..., + help="Path to the Hamilton Python module file.", + exists=True, + file_okay=True, + dir_okay=False, + resolve_path=True, + ), + output: Path | None = typer.Option( + None, + "--output", + "-o", + help="Custom destination path for the generated test file.", + ), +): + """Generate pytest-style unit test stubs for functions in a Hamilton module.""" + out_file = commands.create_tests(module_path=module_path, output_path=output) + _response_handler( + ctx=ctx, + response=Response( + command="create-tests", + success=True, + message=f"Test stubs generated at: {out_file}", + ), + ) + if __name__ == "__main__": cli() diff --git a/hamilton/cli/commands.py b/hamilton/cli/commands.py index d83cd9719..a6c268258 100644 --- a/hamilton/cli/commands.py +++ b/hamilton/cli/commands.py @@ -124,3 +124,8 @@ def version(dr: driver.Driver) -> dict: def view(dr: driver.Driver, output_file_path: Path = Path("dag.png")) -> None: """Display all functions of the instantiated Driver""" dr.display_all_functions(output_file_path) + +def create_tests(module_path: Path, output_path: Path | None = None) -> Path: + """Generate unit test stubs for Hamilton functions in the target module.""" + return logic.create_test_stubs(module_path=module_path, output_file=output_path) + diff --git a/hamilton/cli/logic.py b/hamilton/cli/logic.py index 5a427b791..72ea39362 100644 --- a/hamilton/cli/logic.py +++ b/hamilton/cli/logic.py @@ -336,3 +336,108 @@ def _read_py_context(file_path: Path) -> dict: context[k] = getattr(module, k, None) return context + +def _has_existing_test(test_file_path: Path, test_func_name: str) -> bool: + """Check via AST whether the test file already declares test_func_name.""" + import ast + + if not test_file_path.exists(): + return False + try: + tree = ast.parse(test_file_path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == test_func_name: + return True + except Exception: + return False + return False + + +def generate_test_stub_for_function(func_name: str, func: object, module_name: str) -> str: + """Generate the pytest-style unit test stub for a single Hamilton function.""" + import inspect + + sig = inspect.signature(func) + params = list(sig.parameters.values()) + + param_lines = [] + call_args = [] + has_pandas = False + + for param in params: + annotation_str = "" + if param.annotation != inspect.Parameter.empty: + if hasattr(param.annotation, "__name__"): + annotation_str = param.annotation.__name__ + else: + annotation_str = str(param.annotation).replace("typing.", "") + + if "Series" in annotation_str or "DataFrame" in annotation_str: + has_pandas = True + default_val = "pd.Series([])" if "Series" in annotation_str else "pd.DataFrame({})" + param_lines.append(f" {param.name}: {annotation_str} = {default_val} # TODO: fill in") + elif annotation_str: + param_lines.append(f" {param.name}: {annotation_str} = None # TODO: fill in") + else: + param_lines.append(f" {param.name} = None # TODO: fill in") + call_args.append(param.name) + + call_expr = f"{module_name}.{func_name}({', '.join(call_args)})" + + if has_pandas: + assertion_block = ( + " # TODO: fill in expected\n" + " expected = None\n" + " pd.testing.assert_series_equal(actual, expected)" + ) + else: + assertion_block = ( + " # TODO: fill in expected\n" + " expected = None\n" + " assert actual == expected" + ) + + params_code = "\n".join(param_lines) + if params_code: + params_code += "\n" + + return f"""def test_{func_name}(): +{params_code} actual = {call_expr} +{assertion_block} +""" + + +def create_test_stubs(module_path: Path, output_file: Path | None = None) -> Path: + """Scan a module file and generate unit test stubs for its Hamilton functions.""" + from hamilton import ad_hoc_utils, graph_utils + + module_name = module_path.stem + if output_file is None: + output_file = module_path.parent / f"test_{module_name}.py" + + mod = ad_hoc_utils.module_from_source(module_path.read_text()) + functions = graph_utils.find_functions(mod) + + new_stubs = [] + for func_name, func_obj in functions: + test_func_name = f"test_{func_name}" + if not _has_existing_test(output_file, test_func_name): + new_stubs.append(generate_test_stub_for_function(func_name, func_obj, module_name)) + + if not new_stubs: + return output_file + + if output_file.exists(): + existing_text = output_file.read_text() + spacing = "\n\n" if not existing_text.endswith("\n\n") else "" + output_file.write_text(existing_text + spacing + "\n\n".join(new_stubs) + "\n") + else: + needs_pandas = any("pd." in stub for stub in new_stubs) + header_lines = ["import pytest"] + if needs_pandas: + header_lines.append("import pandas as pd") + header_lines.append(f"import {module_name}\n\n") + header = "\n".join(header_lines) + output_file.write_text(header + "\n\n".join(new_stubs) + "\n") + + return output_file