tetra-rp 0.6.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

tetra_rp/cli/commands/build_utils/scanner.py

@@ -0,0 +1,596 @@
+"""AST scanner for discovering @remote decorated functions and classes."""
+
+import ast
+import importlib
+import logging
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RemoteFunctionMetadata:
+    """Metadata about a @remote decorated function or class."""
+
+    function_name: str
+    module_path: str
+    resource_config_name: str
+    resource_type: str
+    is_async: bool
+    is_class: bool
+    file_path: Path
+    http_method: Optional[str] = None  # HTTP method for LB endpoints: GET, POST, etc.
+    http_path: Optional[str] = None  # HTTP path for LB endpoints: /api/process
+    is_load_balanced: bool = False  # LoadBalancerSlsResource or LiveLoadBalancer
+    is_live_resource: bool = (
+        False  # LiveLoadBalancer (vs deployed LoadBalancerSlsResource)
+    )
+    config_variable: Optional[str] = None  # Variable name like "gpu_config"
+
+
+class RemoteDecoratorScanner:
+    """Scans Python files for @remote decorators and extracts metadata."""
+
+    def __init__(self, project_dir: Path):
+        self.project_dir = project_dir
+        self.py_files: List[Path] = []
+        self.resource_configs: Dict[str, str] = {}  # name -> name
+        self.resource_types: Dict[str, str] = {}  # name -> type
+        self.resource_flags: Dict[str, Dict[str, bool]] = {}  # name -> {flag: bool}
+        self.resource_variables: Dict[str, str] = {}  # name -> variable_name
+
+    def discover_remote_functions(self) -> List[RemoteFunctionMetadata]:
+        """Discover all @remote decorated functions and classes."""
+        functions = []
+
+        # Find all Python files, excluding root-level directories that shouldn't be scanned
+        all_py_files = self.project_dir.rglob("*.py")
+        # Only exclude these directories if they're direct children of project_dir
+        excluded_root_dirs = {".venv", ".flash", ".runpod"}
+        self.py_files = []
+        for f in all_py_files:
+            try:
+                rel_path = f.relative_to(self.project_dir)
+                # Check if first part of path is in excluded_root_dirs
+                if rel_path.parts and rel_path.parts[0] not in excluded_root_dirs:
+                    self.py_files.append(f)
+            except (ValueError, IndexError):
+                # Include files that can't be made relative
+                self.py_files.append(f)
+
+        # First pass: extract all resource configs from all files
+        for py_file in self.py_files:
+            try:
+                content = py_file.read_text(encoding="utf-8")
+                tree = ast.parse(content)
+                self._extract_resource_configs(tree, py_file)
+            except UnicodeDecodeError:
+                logger.debug(f"Skipping non-UTF-8 file: {py_file}")
+            except SyntaxError as e:
+                logger.warning(f"Syntax error in {py_file}: {e}")
+            except Exception as e:
+                logger.debug(f"Failed to parse {py_file}: {e}")
+
+        # Second pass: extract @remote decorated functions
+        for py_file in self.py_files:
+            try:
+                content = py_file.read_text(encoding="utf-8")
+                tree = ast.parse(content)
+                functions.extend(self._extract_remote_functions(tree, py_file))
+            except UnicodeDecodeError:
+                logger.debug(f"Skipping non-UTF-8 file: {py_file}")
+            except SyntaxError as e:
+                logger.warning(f"Syntax error in {py_file}: {e}")
+            except Exception as e:
+                logger.debug(f"Failed to parse {py_file}: {e}")
+
+        return functions
+
+    def _extract_resource_configs(self, tree: ast.AST, py_file: Path) -> None:
+        """Extract resource config variable assignments and determine type flags.
+
+        This method extracts resource configurations and determines is_load_balanced
+        and is_live_resource flags using string-based type matching.
+        """
+        module_path = self._get_module_path(py_file)
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Assign):
+                # Look for assignments like: gpu_config = LiveServerless(...) or api = LiveLoadBalancer(...)
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        variable_name = target.id
+                        config_type = self._get_call_type(node.value)
+
+                        # Accept any class that looks like a resource config (DeployableResource)
+                        if config_type and self._is_resource_config_type(config_type):
+                            # Extract the resource's name parameter (the actual identifier)
+                            # If extraction fails, fall back to variable name
+                            resource_name = self._extract_resource_name(node.value)
+                            if not resource_name:
+                                resource_name = variable_name
+
+                            # Store mapping using the resource's name (or variable name as fallback)
+                            self.resource_configs[resource_name] = resource_name
+                            self.resource_types[resource_name] = config_type
+
+                            # Store variable name for test-mothership config discovery
+                            self.resource_variables[resource_name] = variable_name
+
+                            # Also store variable name mapping for local lookups in same module
+                            var_key = f"{module_path}:{variable_name}"
+                            self.resource_configs[var_key] = resource_name
+                            self.resource_types[var_key] = config_type
+                            self.resource_variables[var_key] = variable_name
+
+                            # Determine boolean flags using string-based type checking
+                            # This is determined by isinstance() at scan time in production,
+                            # but we use string matching for reliability
+                            is_load_balanced = config_type in [
+                                "LoadBalancerSlsResource",
+                                "LiveLoadBalancer",
+                                "CpuLiveLoadBalancer",
+                            ]
+                            is_live_resource = config_type in [
+                                "LiveLoadBalancer",
+                                "CpuLiveLoadBalancer",
+                            ]
+
+                            # Store flags for this resource
+                            self.resource_flags[resource_name] = {
+                                "is_load_balanced": is_load_balanced,
+                                "is_live_resource": is_live_resource,
+                            }
+                            # Also store for variable key
+                            self.resource_flags[var_key] = {
+                                "is_load_balanced": is_load_balanced,
+                                "is_live_resource": is_live_resource,
+                            }
+
+    def _extract_remote_functions(
+        self, tree: ast.AST, py_file: Path
+    ) -> List[RemoteFunctionMetadata]:
+        """Extract @remote decorated functions and classes."""
+        module_path = self._get_module_path(py_file)
+        functions = []
+
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                # Check if this node has @remote decorator
+                remote_decorator = self._find_remote_decorator(node.decorator_list)
+
+                if remote_decorator:
+                    # Extract resource config name from decorator
+                    resource_config_name = self._extract_resource_config_name(
+                        remote_decorator, module_path
+                    )
+
+                    if resource_config_name:
+                        is_async = isinstance(node, ast.AsyncFunctionDef)
+                        is_class = isinstance(node, ast.ClassDef)
+
+                        # Get resource type for this config
+                        resource_type = self._get_resource_type(resource_config_name)
+
+                        # Extract HTTP routing metadata (for LB endpoints)
+                        http_method, http_path = self._extract_http_routing(
+                            remote_decorator
+                        )
+
+                        # Get flags for this resource
+                        flags = self.resource_flags.get(
+                            resource_config_name,
+                            {"is_load_balanced": False, "is_live_resource": False},
+                        )
+
+                        metadata = RemoteFunctionMetadata(
+                            function_name=node.name,
+                            module_path=module_path,
+                            resource_config_name=resource_config_name,
+                            resource_type=resource_type,
+                            is_async=is_async,
+                            is_class=is_class,
+                            file_path=py_file,
+                            http_method=http_method,
+                            http_path=http_path,
+                            is_load_balanced=flags["is_load_balanced"],
+                            is_live_resource=flags["is_live_resource"],
+                            config_variable=self.resource_variables.get(
+                                resource_config_name
+                            ),
+                        )
+                        functions.append(metadata)
+
+        return functions
+
+    def _find_remote_decorator(self, decorators: List[ast.expr]) -> Optional[ast.expr]:
+        """Find @remote decorator in a list of decorators."""
+        for decorator in decorators:
+            # Handle @remote or @remote(...)
+            if isinstance(decorator, ast.Name):
+                if decorator.id == "remote":
+                    return decorator
+            elif isinstance(decorator, ast.Call):
+                if isinstance(decorator.func, ast.Name):
+                    if decorator.func.id == "remote":
+                        return decorator
+                elif isinstance(decorator.func, ast.Attribute):
+                    if decorator.func.attr == "remote":
+                        return decorator
+
+        return None
+
+    def _extract_resource_config_name(
+        self, decorator: ast.expr, module_path: str
+    ) -> Optional[str]:
+        """Extract resource_config name from @remote decorator."""
+        if isinstance(decorator, ast.Name):
+            # @remote without arguments
+            return None
+
+        if isinstance(decorator, ast.Call):
+            # @remote(...) with arguments
+            # Look for resource_config= or first positional arg
+            for keyword in decorator.keywords:
+                if keyword.arg == "resource_config":
+                    return self._extract_name_from_expr(keyword.value, module_path)
+
+            # Try first positional argument
+            if decorator.args:
+                return self._extract_name_from_expr(decorator.args[0], module_path)
+
+        return None
+
+    def _extract_name_from_expr(
+        self, expr: ast.expr, module_path: str
+    ) -> Optional[str]:
+        """Extract config name from an expression (Name or Call).
+
+        Returns the resource's name (from the name= parameter), not the variable name.
+        """
+        if isinstance(expr, ast.Name):
+            # Variable reference: @remote(gpu_config)
+            variable_name = expr.id
+
+            # Try module-scoped lookup first (current module)
+            var_key = f"{module_path}:{variable_name}"
+            if var_key in self.resource_configs:
+                # Return the actual resource name (mapped from variable)
+                return self.resource_configs[var_key]
+
+            # Try simple name lookup
+            if variable_name in self.resource_configs:
+                return self.resource_configs[variable_name]
+
+            # Fall back to the variable name itself (unresolved reference)
+            return variable_name
+
+        elif isinstance(expr, ast.Call):
+            # Direct instantiation: @remote(LiveServerless(name="gpu_config"))
+            # Extract the name= parameter
+            resource_name = self._extract_resource_name(expr)
+            if resource_name:
+                return resource_name
+
+        return None
+
+    def _is_resource_config_type(self, type_name: str) -> bool:
+        """Check if a type represents a ServerlessResource subclass.
+
+        Returns True only if the class can be imported and is a ServerlessResource.
+        """
+        from tetra_rp.core.resources.serverless import ServerlessResource
+
+        try:
+            module = importlib.import_module("tetra_rp")
+            if hasattr(module, type_name):
+                cls = getattr(module, type_name)
+                return isinstance(cls, type) and issubclass(cls, ServerlessResource)
+        except (ImportError, AttributeError, TypeError):
+            pass
+
+        return False
+
+    def _get_call_type(self, expr: ast.expr) -> Optional[str]:
+        """Get the type name of a call expression."""
+        if isinstance(expr, ast.Call):
+            if isinstance(expr.func, ast.Name):
+                return expr.func.id
+            elif isinstance(expr.func, ast.Attribute):
+                return expr.func.attr
+
+        return None
+
+    def _extract_resource_name(self, expr: ast.expr) -> Optional[str]:
+        """Extract the 'name' parameter from a resource config instantiation.
+
+        For example, from LiveServerless(name="01_01_gpu_worker", ...)
+        returns "01_01_gpu_worker".
+        """
+        if isinstance(expr, ast.Call):
+            for keyword in expr.keywords:
+                if keyword.arg == "name":
+                    if isinstance(keyword.value, ast.Constant):
+                        return keyword.value.value
+        return None
+
+    def _get_resource_type(self, resource_config_name: str) -> str:
+        """Get the resource type for a given config name."""
+        if resource_config_name in self.resource_types:
+            return self.resource_types[resource_config_name]
+        # Default to LiveServerless if type not found
+        return "LiveServerless"
+
+    def _sanitize_resource_name(self, name: str) -> str:
+        """Sanitize resource config name for use in filenames.
+
+        Replaces invalid filename characters with underscores and ensures
+        the name starts with a letter or underscore (valid for Python identifiers).
+
+        Args:
+            name: Raw resource config name
+
+        Returns:
+            Sanitized name safe for use in filenames and as Python identifiers
+        """
+        # Replace invalid characters with underscores
+        sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", name)
+
+        # Ensure it starts with a letter or underscore
+        if sanitized and not (sanitized[0].isalpha() or sanitized[0] == "_"):
+            sanitized = f"_{sanitized}"
+
+        return sanitized or "_"
+
+    def _get_module_path(self, py_file: Path) -> str:
+        """Convert file path to module path."""
+        try:
+            # Get relative path from project directory
+            rel_path = py_file.relative_to(self.project_dir)
+
+            # Remove .py extension and convert / to .
+            module = str(rel_path.with_suffix("")).replace("/", ".").replace("\\", ".")
+
+            return module
+        except ValueError:
+            # If relative_to fails, just use filename
+            return py_file.stem
+
+    def _extract_http_routing(
+        self, decorator: ast.expr
+    ) -> tuple[Optional[str], Optional[str]]:
+        """Extract HTTP method and path from @remote decorator.
+
+        Returns:
+            Tuple of (method, path) or (None, None) if not found.
+            method: GET, POST, PUT, DELETE, PATCH
+            path: /api/endpoint routes
+
+        Raises:
+            ValueError: If method is not a valid HTTP verb
+        """
+        if not isinstance(decorator, ast.Call):
+            return None, None
+
+        http_method = None
+        http_path = None
+
+        # Extract keyword arguments: method="POST", path="/api/process"
+        for keyword in decorator.keywords:
+            if keyword.arg == "method":
+                if isinstance(keyword.value, ast.Constant):
+                    http_method = keyword.value.value
+            elif keyword.arg == "path":
+                if isinstance(keyword.value, ast.Constant):
+                    http_path = keyword.value.value
+
+        # Validate HTTP method if provided
+        valid_methods = {"GET", "POST", "PUT", "DELETE", "PATCH"}
+        if http_method is not None and http_method.upper() not in valid_methods:
+            raise ValueError(
+                f"Invalid HTTP method '{http_method}'. Must be one of: {', '.join(valid_methods)}"
+            )
+
+        return http_method, http_path
+
+
+def detect_main_app(
+    project_root: Path, explicit_mothership_exists: bool = False
+) -> Optional[dict]:
+    """Detect main.py FastAPI app and return mothership config.
+
+    Searches for main.py/app.py/server.py and parses AST to find FastAPI app.
+    Only returns config if app has custom routes (not just @remote calls).
+
+    Args:
+        project_root: Root directory of Flash project
+        explicit_mothership_exists: If True, skip auto-detection (explicit config takes precedence)
+
+    Returns:
+        Dict with app metadata: {
+            'file_path': Path,
+            'app_variable': str,
+            'has_routes': bool,
+        }
+        Returns None if no FastAPI app found with custom routes or explicit_mothership_exists is True.
+    """
+    if explicit_mothership_exists:
+        # Explicit mothership config exists, skip auto-detection
+        return None
+    for filename in ["main.py", "app.py", "server.py"]:
+        main_path = project_root / filename
+        if not main_path.exists():
+            continue
+
+        try:
+            content = main_path.read_text(encoding="utf-8")
+            tree = ast.parse(content)
+
+            # Find FastAPI app instantiation
+            for node in ast.walk(tree):
+                if isinstance(node, ast.Assign):
+                    if isinstance(node.value, ast.Call):
+                        call_type = None
+                        if isinstance(node.value.func, ast.Name):
+                            call_type = node.value.func.id
+                        elif isinstance(node.value.func, ast.Attribute):
+                            call_type = node.value.func.attr
+
+                        if call_type == "FastAPI":
+                            app_variable = None
+                            for target in node.targets:
+                                if isinstance(target, ast.Name):
+                                    app_variable = target.id
+                                    break
+
+                            if app_variable:
+                                # Check for custom routes (not just @remote)
+                                has_routes = _has_custom_routes(tree, app_variable)
+
+                                return {
+                                    "file_path": main_path,
+                                    "app_variable": app_variable,
+                                    "has_routes": has_routes,
+                                }
+        except UnicodeDecodeError:
+            logger.debug(f"Skipping non-UTF-8 file: {main_path}")
+        except SyntaxError as e:
+            logger.debug(f"Syntax error in {main_path}: {e}")
+        except Exception as e:
+            logger.debug(f"Failed to parse {main_path}: {e}")
+
+    return None
+
+
+def _has_custom_routes(tree: ast.AST, app_variable: str) -> bool:
+    """Check if FastAPI app has custom routes (beyond @remote).
+
+    Args:
+        tree: AST tree of the file
+        app_variable: Name of the FastAPI app variable
+
+    Returns:
+        True if app has route decorators (app.get, app.post, etc.)
+    """
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            for decorator in node.decorator_list:
+                # Look for app.get(), app.post(), app.put(), etc.
+                if isinstance(decorator, ast.Call):
+                    if isinstance(decorator.func, ast.Attribute):
+                        if (
+                            isinstance(decorator.func.value, ast.Name)
+                            and decorator.func.value.id == app_variable
+                            and decorator.func.attr
+                            in ["get", "post", "put", "delete", "patch"]
+                        ):
+                            return True
+                # Also check for @app.get without parentheses (decorator without Call)
+                elif isinstance(decorator, ast.Attribute):
+                    if (
+                        isinstance(decorator.value, ast.Name)
+                        and decorator.value.id == app_variable
+                        and decorator.attr in ["get", "post", "put", "delete", "patch"]
+                    ):
+                        return True
+
+    return False
+
+
+def detect_explicit_mothership(project_root: Path) -> Optional[Dict]:
+    """Detect explicitly configured mothership resource in mothership.py.
+
+    Parses mothership.py to extract resource configuration.
+
+    Args:
+        project_root: Root directory of Flash project
+
+    Returns:
+        Dict with mothership config if found:
+            {
+                'resource_type': str (e.g., 'CpuLiveLoadBalancer'),
+                'name': str,
+                'workersMin': int,
+                'workersMax': int,
+                'is_explicit': bool,
+            }
+        Returns None if mothership.py doesn't exist or can't be parsed.
+    """
+    mothership_file = project_root / "mothership.py"
+
+    if not mothership_file.exists():
+        return None
+
+    try:
+        content = mothership_file.read_text(encoding="utf-8")
+        tree = ast.parse(content)
+
+        # Look for variable assignment: mothership = SomeResource(...)
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and target.id == "mothership":
+                        # Found mothership variable assignment
+                        if isinstance(node.value, ast.Call):
+                            resource_type = _extract_resource_type(node.value)
+                            kwargs = _extract_call_kwargs(node.value)
+
+                            return {
+                                "resource_type": resource_type,
+                                "name": kwargs.get("name", "mothership"),
+                                "workersMin": kwargs.get("workersMin", 1),
+                                "workersMax": kwargs.get("workersMax", 3),
+                                "is_explicit": True,
+                            }
+
+        return None
+
+    except UnicodeDecodeError:
+        logger.debug(f"Skipping non-UTF-8 file: {mothership_file}")
+        return None
+    except SyntaxError as e:
+        logger.debug(f"Syntax error in mothership.py: {e}")
+        return None
+    except Exception as e:
+        logger.debug(f"Failed to parse mothership.py: {e}")
+        return None
+
+
+def _extract_resource_type(call_node: ast.Call) -> str:
+    """Extract resource type from Call node.
+
+    Args:
+        call_node: AST Call node representing resource instantiation
+
+    Returns:
+        Resource type name (e.g., 'CpuLiveLoadBalancer'), or default if not found
+    """
+    if isinstance(call_node.func, ast.Name):
+        return call_node.func.id
+    elif isinstance(call_node.func, ast.Attribute):
+        return call_node.func.attr
+    return "CpuLiveLoadBalancer"  # Default
+
+
+def _extract_call_kwargs(call_node: ast.Call) -> Dict:
+    """Extract keyword arguments from Call node.
+
+    Args:
+        call_node: AST Call node
+
+    Returns:
+        Dict of keyword arguments with evaluated values (numbers, strings)
+    """
+    kwargs = {}
+    for keyword in call_node.keywords:
+        if keyword.arg:
+            try:
+                # Try to evaluate simple literal values
+                kwargs[keyword.arg] = ast.literal_eval(keyword.value)
+            except (ValueError, SyntaxError, TypeError):
+                # Skip non-literal arguments
+                pass
+    return kwargs