hassle.generate_tests

 1import argparse
 2import os
 3import tokenize
 4from pathlib import Path
 5
 6root = Path(__file__).parent
 7
 8
 9def get_args() -> argparse.Namespace:
10    parser = argparse.ArgumentParser()
11
12    parser.add_argument(
13        "package_name",
14        type=str,
15        default=".",
16        nargs="?",
17        help=""" The name of the package or project to generate tests for,
18        assuming it's a subfolder of your current working directory.
19        Can also be a full path to the package. If nothing is given,
20        the current working directory will be used.""",
21    )
22
23    args = parser.parse_args()
24
25    return args
26
27
28def get_function_names(filepath: Path) -> list[str]:
29    """Returns a list of function names from a .py file."""
30    with filepath.open("r") as file:
31        tokens = list(tokenize.generate_tokens(file.readline))
32    functions = []
33    for i, token in enumerate(tokens):
34        # If token.type is "name" and the preceeding token is "def"
35        if (
36            token.type == 1
37            and tokens[i - 1].type == 1
38            and tokens[i - 1].string == "def"
39        ):
40            functions.append(token.string)
41    return functions
42
43
44def write_placeholders(package_path: Path, pyfile: Path | str, functions: list[str]):
45    """Write placeholder functions to the
46    tests/test_{pyfile} file if they don't already exist.
47    The placeholder functions use the naming convention
48    test__{pyfile.name}__{function_name}
49
50    :param package_path: Path to the package.
51
52    :param pyfile: Path to the pyfile to write placeholders for.
53
54    :param functions: List of functions to generate
55    placehodlers for."""
56    package_name = package_path.stem
57    tests_dir = package_path / "tests"
58    tests_dir.mkdir(parents=True, exist_ok=True)
59    pyfile = Path(pyfile)
60    test_file = tests_dir / f"test_{pyfile.name}"
61    # Makes sure not to overwrite previously written tests
62    # or additional imports.
63    if test_file.exists():
64        content = test_file.read_text() + "\n\n"
65    else:
66        content = f"import pytest\nfrom {package_name} import {pyfile.stem}\n\n\n"
67    for function in functions:
68        test_function = f"def test__{pyfile.stem}__{function}"
69        if test_function not in content and function != "__init__":
70            content += f"{test_function}():\n    ...\n\n\n"
71    test_file.write_text(content)
72    os.system(f"black {tests_dir}")
73    os.system(f"isort {tests_dir}")
74
75
76def generate_test_files(package_path: Path):
77    """Generate test files for all .py files in 'src'
78    directory of 'package_path'."""
79    pyfiles = [
80        file
81        for file in (package_path / "src").rglob("*.py")
82        if file.name != "__init__.py"
83    ]
84    for pyfile in pyfiles:
85        write_placeholders(package_path, pyfile, get_function_names(pyfile))
86
87
88def main(args: argparse.Namespace = None):
89    if not args:
90        args = get_args()
91    package_path = Path(args.package_name).resolve()
92    generate_test_files(package_path)
93
94
95if __name__ == "__main__":
96    main(get_args())
def get_args() -> argparse.Namespace:
10def get_args() -> argparse.Namespace:
11    parser = argparse.ArgumentParser()
12
13    parser.add_argument(
14        "package_name",
15        type=str,
16        default=".",
17        nargs="?",
18        help=""" The name of the package or project to generate tests for,
19        assuming it's a subfolder of your current working directory.
20        Can also be a full path to the package. If nothing is given,
21        the current working directory will be used.""",
22    )
23
24    args = parser.parse_args()
25
26    return args
def get_function_names(filepath: pathlib.Path) -> list[str]:
29def get_function_names(filepath: Path) -> list[str]:
30    """Returns a list of function names from a .py file."""
31    with filepath.open("r") as file:
32        tokens = list(tokenize.generate_tokens(file.readline))
33    functions = []
34    for i, token in enumerate(tokens):
35        # If token.type is "name" and the preceeding token is "def"
36        if (
37            token.type == 1
38            and tokens[i - 1].type == 1
39            and tokens[i - 1].string == "def"
40        ):
41            functions.append(token.string)
42    return functions

Returns a list of function names from a .py file.

def write_placeholders( package_path: pathlib.Path, pyfile: pathlib.Path | str, functions: list[str]):
45def write_placeholders(package_path: Path, pyfile: Path | str, functions: list[str]):
46    """Write placeholder functions to the
47    tests/test_{pyfile} file if they don't already exist.
48    The placeholder functions use the naming convention
49    test__{pyfile.name}__{function_name}
50
51    :param package_path: Path to the package.
52
53    :param pyfile: Path to the pyfile to write placeholders for.
54
55    :param functions: List of functions to generate
56    placehodlers for."""
57    package_name = package_path.stem
58    tests_dir = package_path / "tests"
59    tests_dir.mkdir(parents=True, exist_ok=True)
60    pyfile = Path(pyfile)
61    test_file = tests_dir / f"test_{pyfile.name}"
62    # Makes sure not to overwrite previously written tests
63    # or additional imports.
64    if test_file.exists():
65        content = test_file.read_text() + "\n\n"
66    else:
67        content = f"import pytest\nfrom {package_name} import {pyfile.stem}\n\n\n"
68    for function in functions:
69        test_function = f"def test__{pyfile.stem}__{function}"
70        if test_function not in content and function != "__init__":
71            content += f"{test_function}():\n    ...\n\n\n"
72    test_file.write_text(content)
73    os.system(f"black {tests_dir}")
74    os.system(f"isort {tests_dir}")

Write placeholder functions to the tests/test_{pyfile} file if they don't already exist. The placeholder functions use the naming convention test__{pyfile.name}__{function_name}

Parameters
  • package_path: Path to the package.

  • pyfile: Path to the pyfile to write placeholders for.

  • functions: List of functions to generate placehodlers for.

def generate_test_files(package_path: pathlib.Path):
77def generate_test_files(package_path: Path):
78    """Generate test files for all .py files in 'src'
79    directory of 'package_path'."""
80    pyfiles = [
81        file
82        for file in (package_path / "src").rglob("*.py")
83        if file.name != "__init__.py"
84    ]
85    for pyfile in pyfiles:
86        write_placeholders(package_path, pyfile, get_function_names(pyfile))

Generate test files for all .py files in 'src' directory of 'package_path'.

def main(args: argparse.Namespace = None):
89def main(args: argparse.Namespace = None):
90    if not args:
91        args = get_args()
92    package_path = Path(args.package_name).resolve()
93    generate_test_files(package_path)