review comments

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

Author: akihikokuroda

mellea/helpers/__init__.py

@@ -16,6 +16,7 @@
     wait_for_all_mots,
 )
 from .event_loop_helper import _run_async_in_thread
+from .imports import get_unauthorized_imports
 from .openai_compatible_helpers import (
     chat_completion_delta_merge,
     extract_model_tool_requests,
@@ -36,6 +37,7 @@
     "chat_completion_delta_merge",
     "extract_model_tool_requests",
     "get_current_event_loop",
+    "get_unauthorized_imports",
     "is_vllm_server_with_structured_output",
     "message_to_openai_message",
     "messages_to_docs",

mellea/helpers/imports.py

@@ -0,0 +1,48 @@
+"""Utilities for analyzing Python code imports."""
+
+import ast
+
+
+def get_unauthorized_imports(
+    code: str, allowed_imports: list[str] | None = None
+) -> list[str]:
+    r"""Extract unauthorized imports from Python code.
+
+    Parses Python code and returns a sorted list of top-level modules that are
+    imported but not in the allowed list. Handles both `import X` and `from X import Y`
+    statements, extracting the root module name (e.g., "numpy" from "numpy.random").
+
+    Args:
+        code: Python source code to analyze.
+        allowed_imports: Allowlist of permitted top-level modules. If None, allows all imports.
+
+    Returns:
+        Sorted list of unauthorized module names found in code. Empty list if code
+        has syntax errors or if allowed_imports is None.
+
+    Example:
+        >>> code = "import numpy\nimport os\nimport forbidden_lib"
+        >>> get_unauthorized_imports(code, ["os", "sys"])
+        ['numpy', 'forbidden_lib']
+    """
+    if allowed_imports is None:
+        return []
+
+    unauthorized: set[str] = set()
+    try:
+        tree = ast.parse(code)
+    except (SyntaxError, ValueError):
+        return []
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                module = alias.name.split(".")[0]
+                if module not in allowed_imports:
+                    unauthorized.add(module)
+        elif isinstance(node, ast.ImportFrom) and node.module:
+            module = node.module.split(".")[0]
+            if module not in allowed_imports:
+                unauthorized.add(module)
+
+    return sorted(unauthorized)

mellea/stdlib/requirements/python_tools.py

@@ -56,6 +56,7 @@
 from pathlib import Path
 
 from ...core import Context, Requirement, ValidationResult
+from ...helpers import get_unauthorized_imports
 from .python_reqs import _has_python_code_listing
 
 
@@ -81,41 +82,6 @@ def _extract_code(ctx: Context) -> str | None:
     return None
 
 
-def _get_unauthorized_imports(
-    code: str, allowed_imports: list[str] | None
-) -> list[str]:
-    """Return list of imports in code that are not in allowed_imports.
-
-    Args:
-        code: Python code to analyze
-        allowed_imports: Allowlist of permitted top-level modules (None = allow all)
-
-    Returns:
-        List of unauthorized import module names, or empty list if all allowed.
-    """
-    if allowed_imports is None:
-        return []
-
-    unauthorized = []
-    try:
-        tree = ast.parse(code)
-        for node in ast.walk(tree):
-            if isinstance(node, ast.Import):
-                for alias in node.names:
-                    module_name = alias.name.split(".")[0]
-                    if module_name not in allowed_imports:
-                        unauthorized.append(module_name)
-            elif isinstance(node, ast.ImportFrom):
-                if node.module:
-                    module_name = node.module.split(".")[0]
-                    if module_name not in allowed_imports:
-                        unauthorized.append(module_name)
-    except (SyntaxError, ValueError):
-        pass
-
-    return list(set(unauthorized))
-
-
 def _code_parses(code: str) -> tuple[bool, str | None]:
     """Check if code parses as valid Python.
 
@@ -261,7 +227,7 @@ def validate(ctx: Context) -> ValidationResult:
                 result=False, reason="Could not extract Python code"
             )
 
-        unauthorized = _get_unauthorized_imports(code, allowed_imports)
+        unauthorized = get_unauthorized_imports(code, allowed_imports)
         if unauthorized:
             allowed_str = ", ".join(sorted(set(allowed_imports)))
             return ValidationResult(

mellea/stdlib/tools/interpreter.py

@@ -20,6 +20,7 @@
 from typing import Any
 
 from ...core import MelleaLogger
+from ...helpers import get_unauthorized_imports
 
 logger = MelleaLogger.get_logger()
 
@@ -148,7 +149,7 @@ def execute(self, code: str, timeout: int) -> ExecutionResult:
             )
 
         if self.allowed_imports:
-            unauthorized = _get_unauthorized_imports(code, self.allowed_imports)
+            unauthorized = get_unauthorized_imports(code, self.allowed_imports)
             if unauthorized:
                 return ExecutionResult(
                     success=False,
@@ -185,7 +186,7 @@ def execute(self, code: str, timeout: int) -> ExecutionResult:
             unexpected error occurs.
         """
         if self.allowed_imports:
-            unauthorized = _get_unauthorized_imports(code, self.allowed_imports)
+            unauthorized = get_unauthorized_imports(code, self.allowed_imports)
             if unauthorized:
                 return ExecutionResult(
                     success=False,
@@ -259,7 +260,7 @@ def execute(self, code: str, timeout: int) -> ExecutionResult:
             flag, or a skipped result on import violation or sandbox error.
         """
         if self.allowed_imports:
-            unauthorized = _get_unauthorized_imports(code, self.allowed_imports)
+            unauthorized = get_unauthorized_imports(code, self.allowed_imports)
             if unauthorized:
                 return ExecutionResult(
                     success=False,
@@ -301,37 +302,9 @@ def execute(self, code: str, timeout: int) -> ExecutionResult:
             )
 
 
-def _get_unauthorized_imports(code: str, allowed_imports: list[str]) -> list[str]:
-    """Get list of unauthorized imports used in code."""
-    unauthorized: list[str] = []
-    try:
-        tree = ast.parse(code)
-    except SyntaxError:
-        return unauthorized
-
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Import):
-            for alias in node.names:
-                base_module = alias.name.split(".")[0]
-                if (
-                    base_module not in allowed_imports
-                    and base_module not in unauthorized
-                ):
-                    unauthorized.append(base_module)
-        elif isinstance(node, ast.ImportFrom):
-            if node.module:
-                base_module = node.module.split(".")[0]
-                if (
-                    base_module not in allowed_imports
-                    and base_module not in unauthorized
-                ):
-                    unauthorized.append(base_module)
-    return unauthorized
-
-
 def _check_allowed_imports(code: str, allowed_imports: list[str]) -> bool:
     """Check if code only uses allowed imports."""
-    return len(_get_unauthorized_imports(code, allowed_imports)) == 0
+    return len(get_unauthorized_imports(code, allowed_imports)) == 0
 
 
 def code_interpreter(code: str) -> ExecutionResult: