Add doctest coverage to python scripts

bundle/direct/tools/generate_apitypes.py

@@ -37,7 +37,23 @@ def parse_out_fields(path):
 
 
 def get_schema_fields(schemas):
-    """Get top-level field names for each schema type."""
+    """Get top-level field names for each schema type.
+
+    >>> get_schema_fields({})
+    {}
+    >>> get_schema_fields({"TypeA": {}})
+    {}
+    >>> get_schema_fields({"TypeA": {"fields": {"x": {}}}})
+    {'TypeA': {'x'}}
+    >>> result = get_schema_fields({"A": {"fields": {"x": {}}}, "B": {"fields": {"y": {}, "z": {}}}})
+    >>> result["A"] == {"x"}
+    True
+    >>> result["B"] == {"y", "z"}
+    True
+    >>> result = get_schema_fields({"TypeA": {"fields": {}}, "TypeB": {"fields": {"x": {}}}})
+    >>> "TypeA" not in result and "TypeB" in result
+    True
+    """
     schema_fields = {}
     for name, schema in schemas.items():
         props = schema.get("fields", {})

bundle/direct/tools/generate_resources.py

@@ -111,7 +111,21 @@ def extract(schema, prefix, visited, depth, inherited):
 
 
 def find_inherited_behaviors(schemas, type_name):
-    """Find INPUT_ONLY/OUTPUT_ONLY behaviors from containers that reference type_name."""
+    """Find INPUT_ONLY/OUTPUT_ONLY behaviors from containers that reference type_name.
+
+    >>> find_inherited_behaviors({"A": {"fields": {"f": {"ref": "B", "behaviors": ["INPUT_ONLY"]}}}}, "B")
+    ['INPUT_ONLY']
+    >>> find_inherited_behaviors({"A": {"fields": {"f": {"ref": "B", "behaviors": ["OUTPUT_ONLY"]}}}}, "B")
+    ['OUTPUT_ONLY']
+    >>> find_inherited_behaviors({"A": {"fields": {"f": {"ref": "B", "behaviors": ["INPUT_ONLY", "OUTPUT_ONLY"]}}}}, "B")
+    ['INPUT_ONLY', 'OUTPUT_ONLY']
+    >>> find_inherited_behaviors({"A": {"fields": {"f": {"ref": "B"}}}}, "B")
+    []
+    >>> find_inherited_behaviors({}, "B")
+    []
+    >>> find_inherited_behaviors({"A": {"fields": {}}}, "B")
+    []
+    """
     inherited = []
     for container_schema in schemas.values():
         for field_prop in container_schema.get("fields", {}).values():
@@ -126,7 +140,21 @@ def find_inherited_behaviors(schemas, type_name):
 
 
 def filter_prefixes(fields):
-    """Remove fields that are children of other fields in the list."""
+    """Remove fields that are children of other fields in the list.
+
+    >>> filter_prefixes([("a.b.c", "X"), ("a", "Y"), ("a.b", "Z")])
+    [('a', 'Y')]
+    >>> filter_prefixes([("x", "A"), ("y", "B")])
+    [('x', 'A'), ('y', 'B')]
+    >>> filter_prefixes([])
+    []
+    >>> filter_prefixes([("field", "IMMUTABLE")])
+    [('field', 'IMMUTABLE')]
+    >>> filter_prefixes([("config", "A"), ("config.nested", "B"), ("config.nested.deep", "C")])
+    [('config', 'A')]
+    >>> filter_prefixes([("a", "X"), ("b.c", "Y")])
+    [('a', 'X'), ('b.c', 'Y')]
+    """
     result = []
     for field, behavior in sorted(fields):
         if not any(field.startswith(f + ".") for f, _ in result):
@@ -135,7 +163,23 @@ def filter_prefixes(fields):
 
 
 def write_field_group(lines, header, fields):
-    """Write a group of fields with field and reason, grouped by behavior."""
+    """Write a group of fields with field and reason, grouped by behavior.
+
+    >>> lines = []
+    >>> write_field_group(lines, "test_group", [("f1", "IMMUTABLE")])
+    >>> len(lines)
+    3
+    >>> "test_group" in lines[0]
+    True
+    >>> lines = []
+    >>> write_field_group(lines, "multi", [("a", "OUTPUT_ONLY"), ("b", "INPUT_ONLY")])
+    >>> len(lines)
+    6
+    >>> any("spec:input_only" in line for line in lines)
+    True
+    >>> any("spec:output_only" in line for line in lines)
+    True
+    """
     lines.append(f"\n    {header}:")
     # Group by behavior
     by_behavior = {}
@@ -153,7 +197,27 @@ def write_field_group(lines, header, fields):
 
 
 def generate(resource_behaviors):
-    """Generate resources.yml."""
+    """Generate resources.yml.
+
+    >>> result = generate({"res1": {"f1": ["OUTPUT_ONLY"]}})
+    >>> "res1:" in result
+    True
+    >>> "ignore_remote_changes:" in result
+    True
+    >>> "spec:output_only" in result
+    True
+    >>> result = generate({})
+    >>> "resources:" in result
+    True
+    >>> result = generate({"res": {}})
+    >>> "no api field behaviors" in result
+    True
+    >>> result = generate({"res": {"field": ["IMMUTABLE", "OUTPUT_ONLY"]}})
+    >>> "recreate_on_changes:" in result
+    True
+    >>> "ignore_remote_changes:" in result
+    True
+    """
     lines = [
         """# Generated, do not edit. API field behaviors from OpenAPI schema.
 #

tools/bench_parse.py

@@ -14,6 +14,26 @@ def parse_key_values(text):
 
     >>> parse_key_values("wall=10.316  ru_utime=19.207  ru_stime=0.505  ru_maxrss=573079552")
     {'wall': 10.316, 'ru_utime': 19.207, 'ru_stime': 0.505, 'ru_maxrss': 573079552.0}
+
+    Empty string returns empty dict:
+
+    >>> parse_key_values("")
+    {}
+
+    Single pair:
+
+    >>> parse_key_values("key=42")
+    {'key': 42.0}
+
+    Non-numeric values are kept as strings:
+
+    >>> parse_key_values("name=test  count=5")
+    {'name': 'test', 'count': 5.0}
+
+    Pairs without equals sign are skipped:
+
+    >>> parse_key_values("valid=1  invalid  also_valid=2")
+    {'valid': 1.0, 'also_valid': 2.0}
     """
     result = {}
     for kv_pair in text.split():
@@ -53,7 +73,24 @@ def parse_bench_output(file_path):
 
 
 def calculate_means(results):
-    """Calculate mean values for each metric."""
+    """Calculate mean values for each metric.
+
+    >>> calculate_means({"test1": {"wall": [10, 20]}})
+    {'test1': {'wall': 15}}
+
+    >>> calculate_means({"test1": {"wall": [1.5, 2.5, 3.0]}})  # doctest: +ELLIPSIS
+    {'test1': {'wall': 2.33...}}
+
+    Multiple tests and metrics:
+
+    >>> sorted(calculate_means({"t1": {"m1": [2, 4], "m2": [100]}, "t2": {"m1": [3]}}).items())
+    [('t1', {'m1': 3, 'm2': 100}), ('t2', {'m1': 3})]
+
+    Empty metrics list returns zero:
+
+    >>> calculate_means({"test": {"metric": []}})
+    {'test': {'metric': 0}}
+    """
     means = {}
     for test_name, metrics in results.items():
         means[test_name] = {metric: statistics.mean(values) if values else 0 for metric, values in metrics.items()}

tools/bump_mlops_stacks.py

@@ -38,7 +38,23 @@
 
 
 def render_path(rel_path, config):
-    """Map a path under the upstream template/ dir to its path in the rendered project."""
+    """Map a path under the upstream template/ dir to its path in the rendered project.
+
+    >>> render_path(Path("template/{{.input_root_dir}}/job.yml"), {"input_root_dir": "src", "input_project_name": "myproj"})
+    'template/src/job.yml'
+
+    >>> render_path(Path("template/{{.input_project_name}}/config.json.tmpl"), {"input_root_dir": "src", "input_project_name": "myproj"})
+    'template/myproj/config.json'
+
+    Handles both vendor and upstream project directory patterns:
+    >>> cfg = {"input_root_dir": ".", "input_project_name": "proj"}
+    >>> render_path(Path("template/{{.input_root_dir}}/{{template `project_name_alphanumeric_underscore` .}}/file"), cfg)
+    'template/./proj/file'
+
+    Removes .tmpl suffix:
+    >>> render_path(Path("template/README.md.tmpl"), {"input_root_dir": ".", "input_project_name": "p"})
+    'template/README.md'
+    """
     subs = {
         PROJECT_ROOT_SEG: config["input_root_dir"],
         VENDORED_PROJECT_DIR: config["input_project_name"],
@@ -85,6 +101,17 @@ def keep_set(clone_dir, out_dir, config):
 
 
 def vendored_path(rel_path):
+    """Replace upstream project directory pattern with vendored pattern.
+
+    >>> vendored_path(Path("template/{{template `project_name_alphanumeric_underscore` .}}/config"))
+    PosixPath('template/{{.input_project_name}}/config')
+
+    >>> vendored_path(Path("library/helpers.py"))
+    PosixPath('library/helpers.py')
+
+    >>> vendored_path(Path("a/{{template `project_name_alphanumeric_underscore` .}}/b/{{template `project_name_alphanumeric_underscore` .}}/c"))
+    PosixPath('a/{{.input_project_name}}/b/{{.input_project_name}}/c')
+    """
     return Path(*[VENDORED_PROJECT_DIR if p == UPSTREAM_PROJECT_DIR else p for p in rel_path.parts])
 
 

tools/check_deadcode.py

@@ -47,6 +47,39 @@
 ALLOW_COMMENT = "//deadcode:allow"
 
 
+def should_exclude_line(line, excluded_dirs):
+    """Check if a deadcode output line refers to an excluded directory.
+
+    >>> should_exclude_line("libs/gorules/myrule.go:10:5: func", ["libs/gorules/"])
+    True
+    >>> should_exclude_line("bundle/internal/tf/schema/gen.go:10:5: func", ["bundle/internal/tf/schema/"])
+    True
+    >>> should_exclude_line("cmd/bundle/deploy.go:10:5: func", ["libs/gorules/"])
+    False
+    >>> should_exclude_line("bundle/internal/tf/schema/gen.go:10:5: func", ["libs/gorules/"])
+    False
+    """
+    return any(line.startswith(d) or ("/" + d) in line for d in excluded_dirs)
+
+
+def parse_deadcode_line(line):
+    """Parse a deadcode output line into (filepath, lineno) or return None if unparseable.
+
+    Typical deadcode format: path/to/file.go:123:45: message
+    >>> parse_deadcode_line("cmd/main.go:42:3: func Foo")
+    ('cmd/main.go', 42)
+    >>> parse_deadcode_line("libs/util/helper.go:1:0: func Helper")
+    ('libs/util/helper.go', 1)
+    >>> parse_deadcode_line("invalid line format")
+
+    >>> parse_deadcode_line("")
+    """
+    match = re.match(r"(.+?):(\d+):\d+:", line)
+    if not match:
+        return None
+    return (match.group(1), int(match.group(2)))
+
+
 def main():
     result = subprocess.run(
         ["go", "tool", "-modfile=tools/go.mod", "deadcode", "-test", "./..."],
@@ -67,16 +100,15 @@ def main():
     violations = []
 
     for line in lines:
-        if any(line.startswith(d) or ("/" + d) in line for d in EXCLUDED_DIRS):
+        if should_exclude_line(line, EXCLUDED_DIRS):
             continue
 
-        match = re.match(r"(.+?):(\d+):\d+:", line)
-        if not match:
+        parsed = parse_deadcode_line(line)
+        if not parsed:
             violations.append(line)
             continue
 
-        filepath = match.group(1)
-        lineno = int(match.group(2))
+        filepath, lineno = parsed
 
         try:
             with open(filepath) as f: