dss/.dss/doc-sync/generators/api_extractor.py

#!/usr/bin/env python3
"""
API Extractor

Extract FastAPI route definitions from server.py files.
"""

import ast
import re
from pathlib import Path
from typing import Dict, List, Any, Optional
import logging

from .base_generator import DocGenerator

logger = logging.getLogger(__name__)


class APIExtractor(DocGenerator):
    """
    Extract FastAPI endpoints from server.py files.

    Extracts:
    - Route paths and HTTP methods
    - Function names and docstrings
    - Route parameters
    - Response models
    """

    def extract(self, source_path: Path) -> Dict[str, Any]:
        """
        Extract API endpoints from FastAPI server file.

        Args:
            source_path: Path to server.py file

        Returns:
            Dictionary with extracted endpoint data
        """
        logger.info(f"Extracting API endpoints from {source_path}")

        with open(source_path, 'r') as f:
            source_code = f.read()

        tree = ast.parse(source_code)

        endpoints = []
        app_mounts = []

        # Find @app.get, @app.post, etc. decorators
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                endpoint = self._extract_endpoint(node, source_code)
                if endpoint:
                    endpoints.append(endpoint)

            # Find app.mount() calls for static files
            if isinstance(node, ast.Expr):
                mount = self._extract_mount(node)
                if mount:
                    app_mounts.append(mount)

        return {
            "source_file": str(source_path),
            "endpoints": endpoints,
            "mounts": app_mounts,
            "total_endpoints": len(endpoints),
            "total_mounts": len(app_mounts)
        }

    def _extract_endpoint(
        self,
        func_node: ast.FunctionDef,
        source_code: str
    ) -> Optional[Dict[str, Any]]:
        """
        Extract endpoint information from function with decorator.

        Args:
            func_node: AST function definition node
            source_code: Full source code (for extracting decorator args)

        Returns:
            Endpoint data or None
        """
        for decorator in func_node.decorator_list:
            # Check if decorator is app.get, app.post, etc.
            if isinstance(decorator, ast.Call):
                if isinstance(decorator.func, ast.Attribute):
                    # app.get("/path")
                    if decorator.func.attr in ['get', 'post', 'put', 'delete', 'patch']:
                        method = decorator.func.attr.upper()
                        path = self._extract_route_path(decorator)

                        return {
                            "path": path,
                            "method": method,
                            "function": func_node.name,
                            "docstring": ast.get_docstring(func_node),
                            "parameters": self._extract_parameters(func_node),
                            "line_number": func_node.lineno
                        }

        return None

    def _extract_route_path(self, decorator: ast.Call) -> str:
        """
        Extract route path from decorator arguments.

        Args:
            decorator: AST Call node for decorator

        Returns:
            Route path string
        """
        if decorator.args:
            first_arg = decorator.args[0]
            if isinstance(first_arg, ast.Constant):
                return first_arg.value
            elif isinstance(first_arg, ast.Str):  # Python 3.7 compatibility
                return first_arg.s

        return "/"

    def _extract_parameters(self, func_node: ast.FunctionDef) -> List[Dict[str, str]]:
        """
        Extract function parameters.

        Args:
            func_node: AST function definition node

        Returns:
            List of parameter dictionaries
        """
        params = []

        for arg in func_node.args.args:
            param = {"name": arg.arg}

            # Extract type annotation if present
            if arg.annotation:
                param["type"] = ast.unparse(arg.annotation) if hasattr(ast, 'unparse') else str(arg.annotation)

            params.append(param)

        return params

    def _extract_mount(self, expr_node: ast.Expr) -> Optional[Dict[str, Any]]:
        """
        Extract app.mount() call for static files.

        Args:
            expr_node: AST expression node

        Returns:
            Mount data or None
        """
        if isinstance(expr_node.value, ast.Call):
            call = expr_node.value

            # Check if it's app.mount()
            if isinstance(call.func, ast.Attribute):
                if call.func.attr == 'mount' and len(call.args) >= 2:
                    path_arg = call.args[0]
                    mount_path = None

                    if isinstance(path_arg, ast.Constant):
                        mount_path = path_arg.value
                    elif isinstance(path_arg, ast.Str):
                        mount_path = path_arg.s

                    if mount_path:
                        return {
                            "path": mount_path,
                            "type": "StaticFiles",
                            "line_number": expr_node.lineno
                        }

        return None

    def transform(self, extracted_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Transform extracted API data to .knowledge/dss-architecture.json schema.

        Args:
            extracted_data: Raw extracted endpoint data

        Returns:
            Transformed data for knowledge base
        """
        # Read existing architecture.json
        target_path = self.project_root / ".knowledge" / "dss-architecture.json"
        existing = self.read_existing_target(target_path)

        if not existing:
            # Create new structure
            existing = {
                "$schema": "dss-knowledge-v1",
                "type": "architecture",
                "version": "1.0.0",
                "last_updated": None,
                "modules": []
            }

        # Ensure modules list exists
        if "modules" not in existing:
            existing["modules"] = []

        # Create REST API module data
        rest_api_module = {
            "name": "rest_api",
            "path": extracted_data["source_file"],
            "purpose": "FastAPI server providing REST API and static file serving",
            "port": 3456,
            "endpoints": extracted_data["endpoints"],
            "mounts": extracted_data["mounts"],
            "total_endpoints": extracted_data["total_endpoints"]
        }

        # Update or append REST API module
        rest_api_index = next(
            (i for i, m in enumerate(existing["modules"]) if m.get("name") == "rest_api"),
            None
        )

        if rest_api_index is not None:
            existing["modules"][rest_api_index] = rest_api_module
        else:
            existing["modules"].append(rest_api_module)

        existing["last_updated"] = self.metadata["generated_at"]

        return existing