Merge branch 'autogen' into test/test-utils-load_resource

Author: chalmerlowe

scripts/microgenerator/generate.py

@@ -24,6 +24,7 @@
 
 import ast
 import os
+import argparse
 import glob
 import logging
 import re
@@ -492,3 +493,259 @@ def analyze_source_files(
 
     return parsed_data, all_imports, all_types, request_arg_schema
 
+
+# =============================================================================
+# Section 3: Code Generation
+# =============================================================================
+
+
+def _generate_import_statement(
+    context: List[Dict[str, Any]], key: str, package: str
+) -> str:
+    """Generates a formatted import statement from a list of context dictionaries.
+
+    Args:
+        context: A list of dictionaries containing the data.
+        key: The key to extract from each dictionary in the context.
+        package: The base import package (e.g., "google.cloud.bigquery_v2.services").
+
+    Returns:
+        A formatted, multi-line import statement string.
+    """
+    names = sorted(list(set([item[key] for item in context])))
+    names_str = ",\n    ".join(names)
+    return f"from {package} import (\n    {names_str}\n)"
+
+
+def _get_request_class_name(method_name: str, config: Dict[str, Any]) -> str:
+    """Gets the inferred request class name, applying overrides from config."""
+    inferred_request_name = name_utils.method_to_request_class_name(method_name)
+    method_overrides = config.get("filter", {}).get("methods", {}).get("overrides", {})
+    if method_name in method_overrides:
+        return method_overrides[method_name].get(
+            "request_class_name", inferred_request_name
+        )
+    return inferred_request_name
+
+
+def _find_fq_request_name(
+    request_name: str, request_arg_schema: Dict[str, List[str]]
+) -> str:
+    """Finds the fully qualified request name in the schema."""
+    for key in request_arg_schema.keys():
+        if key.endswith(f".{request_name}"):
+            return key
+    return ""
+
+
+def generate_code(config: Dict[str, Any], analysis_results: tuple) -> None:
+    """
+    Generates source code files using Jinja2 templates.
+    """
+    data, all_imports, all_types, request_arg_schema = analysis_results
+    project_root = config["project_root"]
+    config_dir = config["config_dir"]
+
+    templates_config = config.get("templates", [])
+    for item in templates_config:
+        template_path = str(Path(config_dir) / item["template"])
+        output_path = str(Path(project_root) / item["output"])
+
+        template = utils.load_template(template_path)
+        methods_context = []
+        for class_name, methods in data.items():
+            for method_name, method_info in methods.items():
+                context = {
+                    "name": method_name,
+                    "class_name": class_name,
+                    "return_type": method_info["return_type"],
+                }
+
+                request_name = _get_request_class_name(method_name, config)
+                fq_request_name = _find_fq_request_name(
+                    request_name, request_arg_schema
+                )
+
+                if fq_request_name:
+                    context["request_class_full_name"] = fq_request_name
+                    context["request_id_args"] = request_arg_schema[fq_request_name]
+
+                methods_context.append(context)
+
+        # Prepare imports for the template
+        services_context = []
+        client_class_names = sorted(
+            list(set([m["class_name"] for m in methods_context]))
+        )
+
+        for class_name in client_class_names:
+            service_name_cluster = name_utils.generate_service_names(class_name)
+            services_context.append(service_name_cluster)
+
+        # Also need to update methods_context to include the service_name and module_name
+        # so the template knows which client to use for each method.
+        class_to_service_map = {s["service_client_class"]: s for s in services_context}
+        for method in methods_context:
+            service_info = class_to_service_map.get(method["class_name"])
+            if service_info:
+                method["service_name"] = service_info["service_name"]
+                method["service_module_name"] = service_info["service_module_name"]
+
+        # Prepare new imports
+        service_imports = [
+            _generate_import_statement(
+                services_context,
+                "service_module_name",
+                "google.cloud.bigquery_v2.services",
+            )
+        ]
+
+        # Prepare type imports
+        type_imports = [
+            _generate_import_statement(
+                services_context, "service_name", "google.cloud.bigquery_v2.types"
+            )
+        ]
+
+        final_code = template.render(
+            service_name=config.get("service_name"),
+            methods=methods_context,
+            services=services_context,
+            service_imports=service_imports,
+            type_imports=type_imports,
+            request_arg_schema=request_arg_schema,
+        )
+
+        utils.write_code_to_file(output_path, final_code)
+
+
+# =============================================================================
+# Section 4: Main Execution
+# =============================================================================
+
+
+def setup_config_and_paths(config_path: str) -> Dict[str, Any]:
+    """Loads the configuration and sets up necessary paths.
+
+    Args:
+        config_path: The path to the YAML configuration file.
+
+    Returns:
+        A dictionary containing the loaded configuration and paths.
+    """
+
+    def find_project_root(start_path: str, markers: list[str]) -> str | None:
+        """Finds the project root by searching upwards for a marker."""
+        current_path = os.path.abspath(start_path)
+        while True:
+            for marker in markers:
+                if os.path.exists(os.path.join(current_path, marker)):
+                    return current_path
+            parent_path = os.path.dirname(current_path)
+            if parent_path == current_path:  # Filesystem root
+                return None
+            current_path = parent_path
+
+    # Load configuration from the YAML file.
+    config = utils.load_config(config_path)
+
+    # Determine the project root.
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    project_root = find_project_root(script_dir, ["setup.py", ".git"])
+    if not project_root:
+        project_root = os.getcwd()  # Fallback to current directory
+
+    # Set paths in the config dictionary.
+    config["project_root"] = project_root
+    config["config_dir"] = os.path.dirname(os.path.abspath(config_path))
+
+    return config
+
+
+def _execute_post_processing(config: Dict[str, Any]):
+    """
+    Executes post-processing steps, such as patching existing files.
+    """
+    project_root = config["project_root"]
+    post_processing_jobs = config.get("post_processing_templates", [])
+
+    for job in post_processing_jobs:
+        template_path = os.path.join(config["config_dir"], job["template"])
+        target_file_path = os.path.join(project_root, job["target_file"])
+
+        if not os.path.exists(target_file_path):
+            logging.warning(
+                f"Target file {target_file_path} not found, skipping post-processing job."
+            )
+            continue
+
+        # Read the target file
+        with open(target_file_path, "r") as f:
+            lines = f.readlines()
+
+        # --- Extract existing imports and __all__ members ---
+        imports = []
+        all_list = []
+        all_start_index = -1
+        all_end_index = -1
+
+        for i, line in enumerate(lines):
+            if line.strip().startswith("from ."):
+                imports.append(line.strip())
+            if line.strip() == "__all__ = (":
+                all_start_index = i
+            if all_start_index != -1 and line.strip() == ")":
+                all_end_index = i
+
+        if all_start_index != -1 and all_end_index != -1:
+            for i in range(all_start_index + 1, all_end_index):
+                member = lines[i].strip().replace('"', "").replace(",", "")
+                if member:
+                    all_list.append(member)
+
+        # --- Add new items and sort ---
+        for new_import in job.get("add_imports", []):
+            if new_import not in imports:
+                imports.append(new_import)
+        imports.sort()
+        imports = [f"{imp}\n" for imp in imports]  # re-add newlines
+
+        for new_member in job.get("add_to_all", []):
+            if new_member not in all_list:
+                all_list.append(new_member)
+        all_list.sort()
+
+        # --- Render the new file content ---
+        template = utils.load_template(template_path)
+        new_content = template.render(
+            imports=imports,
+            all_list=all_list,
+        )
+
+        # --- Overwrite the target file ---
+        with open(target_file_path, "w") as f:
+            f.write(new_content)
+
+        logging.info(f"Successfully post-processed and overwrote {target_file_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="A generic Python code generator for clients."
+    )
+    parser.add_argument("config", help="Path to the YAML configuration file.")
+    args = parser.parse_args()
+
+    # Load config and set up paths.
+    config = setup_config_and_paths(args.config)
+
+    # Analyze the source code.
+    analysis_results = analyze_source_files(config)
+
+    # Generate the new client code.
+    generate_code(config, analysis_results)
+
+    # Run post-processing steps.
+    _execute_post_processing(config)
+
+    # TODO: Ensure blacken gets called on the generated source files as a final step.