Add imports parameter to django_shell tool and refactor code parsing (#27)

Author: joshuadavidthomas

CHANGELOG.md

@@ -18,6 +18,10 @@ and this project attempts to adhere to [Semantic Versioning](https://semver.org/
 
 ## [Unreleased]
 
+### Added
+
+- Optional `imports` parameter to `django_shell` tool for providing import statements separately from main code execution.
+
 ## [0.7.0]
 
 ### Added

src/mcp_django_shell/code.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import ast
+from typing import Any
+from typing import Literal
+
+
+def filter_existing_imports(imports: str, globals_dict: dict[str, Any]) -> str:
+    tree = ast.parse(imports)
+
+    if not all(isinstance(stmt, ast.Import | ast.ImportFrom) for stmt in tree.body):
+        raise ValueError("Input must contain only import statements")
+
+    filtered_statements: list[tuple[ast.Import | ast.ImportFrom, list[ast.alias]]] = []
+
+    for stmt in tree.body:
+        if isinstance(stmt, ast.Import):
+            needed = [
+                alias
+                for alias in stmt.names
+                if (alias.asname or alias.name.split(".")[0]) not in globals_dict
+            ]
+            if needed:
+                filtered_statements.append((stmt, needed))
+
+        elif isinstance(stmt, ast.ImportFrom):
+            if stmt.names[0].name == "*":
+                # Star imports - always include (can't easily determine what's imported)
+                filtered_statements.append((stmt, stmt.names))
+            else:
+                needed = [
+                    alias
+                    for alias in stmt.names
+                    if (alias.asname or alias.name) not in globals_dict
+                ]
+                if needed:
+                    filtered_statements.append((stmt, needed))
+
+    filtered_lines: list[str] = []
+
+    for stmt, names in filtered_statements:
+        import_parts: list[str] = []
+
+        if isinstance(stmt, ast.ImportFrom):
+            module = stmt.module or ""
+            level = "." * stmt.level if stmt.level else ""
+            import_parts.append(f"from {level}{module}")
+
+        import_parts.append("import")
+
+        name_parts = [
+            f"{alias.name} as {alias.asname}" if alias.asname else alias.name
+            for alias in names
+        ]
+        import_parts.append(", ".join(name_parts))
+
+        filtered_lines.append(" ".join(import_parts))
+
+    return "\n".join(filtered_lines)
+
+
+def parse_code(code: str) -> tuple[str, str, Literal["expression", "statement"]]:
+    """Determine how code should be executed.
+
+    Returns:
+        A tuple (main_code, setup_code, code_type), where:
+        - main_code: The code to evaluate (expression) or execute (statement)
+        - setup_code: Lines to execute before evaluating expressions (empty for statements)
+        - code_type: "expression" or "statement"
+    """
+
+    def can_eval(code: str) -> bool:
+        try:
+            compile(code, "<stdin>", "eval")
+            return True
+        except SyntaxError:
+            return False
+
+    if can_eval(code):
+        return code, "", "expression"
+
+    lines = code.strip().splitlines()
+    last_line = lines[-1] if lines else ""
+
+    if can_eval(last_line):
+        setup_code = "\n".join(lines[:-1])
+        return last_line, setup_code, "expression"
+
+    return code, "", "statement"

src/mcp_django_shell/server.py

@@ -8,6 +8,8 @@
 from fastmcp import FastMCP
 from mcp.types import ToolAnnotations
 
+from .code import filter_existing_imports
+from .code import parse_code
 from .output import DjangoShellOutput
 from .output import ErrorOutput
 from .resources import AppResource
@@ -55,8 +57,12 @@
     ),
 )
 async def django_shell(
-    code: Annotated[str, "Python code to be executed inside the Django shell session"],
     ctx: Context,
+    code: Annotated[str, "Python code to be executed inside the Django shell session"],
+    imports: Annotated[
+        str | None,
+        "Optional import statements to execute before running the main code. Should contain all necessary imports for the code to run successfully, such as 'from django.contrib.auth.models import User\\nfrom myapp.models import MyModel'",
+    ] = None,
 ) -> DjangoShellOutput:
     """Execute Python code in a stateful Django shell session.
 
@@ -72,17 +78,28 @@ async def django_shell(
     `.filter()` and `.get()` rather than their async counterparts (`.afilter()`, `.aget()`).
     """
     logger.info(
-        "django_shell tool called - request_id: %s, client_id: %s, code: %s",
+        "django_shell tool called - request_id: %s, client_id: %s, code: %s, imports: %s",
         ctx.request_id,
         ctx.client_id or "unknown",
         (code[:100] + "..." if len(code) > 100 else code).replace("\n", "\\n"),
+        (imports[:50] + "..." if imports and len(imports) > 50 else imports or "None"),
     )
     logger.debug(
         "Full code for django_shell - request_id: %s: %s", ctx.request_id, code
     )
+    if imports:
+        logger.debug(
+            "Imports for django_shell - request_id: %s: %s", ctx.request_id, imports
+        )
+
+        filtered_imports = filter_existing_imports(imports, shell.globals)
+        if filtered_imports.strip():
+            code = f"{filtered_imports}\n{code}"
+
+    parsed_code, setup, code_type = parse_code(code)
 
     try:
-        result = await shell.execute(code)
+        result = await shell.execute(parsed_code, setup, code_type)
         output = DjangoShellOutput.from_result(result)
 
         logger.debug(

src/mcp_django_shell/shell.py

@@ -42,7 +42,9 @@ def reset(self):
         self.globals = {}
         self.history = []
 
-    async def execute(self, code: str) -> Result:
+    async def execute(
+        self, code: str, setup: str, code_type: Literal["expression", "statement"]
+    ) -> Result:
         """Execute Python code in the Django shell context (async wrapper).
 
         This async wrapper enables use from FastMCP and other async contexts.
@@ -54,9 +56,11 @@ async def execute(self, code: str) -> Result:
         errors.
         """
 
-        return await sync_to_async(self._execute)(code)
+        return await sync_to_async(self._execute)(code, setup, code_type)
 
-    def _execute(self, code: str) -> Result:
+    def _execute(
+        self, code: str, setup: str, code_type: Literal["expression", "statement"]
+    ) -> Result:
         """Execute Python code in the Django shell context (synchronous).
 
         Attempts to evaluate code as an expression first (returning a value),
@@ -65,6 +69,7 @@ def _execute(self, code: str) -> Result:
         Note: This synchronous method contains the actual execution logic.
         Use `execute()` for async contexts or `_execute()` for sync/testing.
         """
+
         code_preview = (code[:100] + "..." if len(code) > 100 else code).replace(
             "\n", "\\n"
         )
@@ -75,8 +80,6 @@ def _execute(self, code: str) -> Result:
 
         with redirect_stdout(stdout), redirect_stderr(stderr):
             try:
-                code, setup, code_type = parse_code(code)
-
                 logger.debug(
                     "Execution type: %s, has setup: %s", code_type, bool(setup)
                 )
@@ -85,7 +88,6 @@ def _execute(self, code: str) -> Result:
                     code[:200] + "..." if len(code) > 200 else code,
                 )
 
-                # Execute setup, if any (only applicable to expressions)
                 if setup:
                     logger.debug(
                         "Setup code: %s",
@@ -146,36 +148,6 @@ def save_result(self, result: Result) -> Result:
         return result
 
 
-def parse_code(code: str) -> tuple[str, str, Literal["expression", "statement"]]:
-    """Determine how code should be executed.
-
-    Returns:
-        A tuple (main_code, setup_code, code_type), where:
-        - main_code: The code to evaluate (expression) or execute (statement)
-        - setup_code: Lines to execute before evaluating expressions (empty for statements)
-        - code_type: "expression" or "statement"
-    """
-
-    def can_eval(code: str) -> bool:
-        try:
-            compile(code, "<stdin>", "eval")
-            return True
-        except SyntaxError:
-            return False
-
-    if can_eval(code):
-        return code, "", "expression"
-
-    lines = code.strip().splitlines()
-    last_line = lines[-1] if lines else ""
-
-    if can_eval(last_line):
-        setup_code = "\n".join(lines[:-1])
-        return last_line, setup_code, "expression"
-
-    return code, "", "statement"
-
-
 @dataclass
 class ExpressionResult:
     code: str

tests/test_code.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import pytest
+
+from mcp_django_shell.code import filter_existing_imports
+from mcp_django_shell.code import parse_code
+
+
+def test_parse_code_single_expression():
+    code, setup, code_type = parse_code("2 + 2")
+
+    assert code == "2 + 2"
+    assert setup == ""
+    assert code_type == "expression"
+
+
+def test_parse_code_single_statement():
+    code, setup, code_type = parse_code("x = 5")
+
+    assert code == "x = 5"
+    assert setup == ""
+    assert code_type == "statement"
+
+
+def test_parse_code_multiline_with_expression_basic():
+    code, setup, code_type = parse_code("x = 5\ny = 10\nx + y")
+
+    assert code == "x + y"
+    assert setup == "x = 5\ny = 10"
+    assert code_type == "expression"
+
+
+def test_parse_code_multiline_statement_only():
+    code, setup, code_type = parse_code("x = 5\ny = 10\nz = x + y")
+
+    assert code == "x = 5\ny = 10\nz = x + y"
+    assert setup == ""
+    assert code_type == "statement"
+
+
+def test_parse_code_empty_code():
+    code, setup, code_type = parse_code("")
+
+    assert code == ""
+    assert setup == ""
+    assert code_type == "statement"
+
+
+def test_parse_code_whitespace_only():
+    code, setup, code_type = parse_code("   \n  \t  ")
+
+    assert code == "   \n  \t  "
+    assert setup == ""
+    assert code_type == "statement"
+
+
+def test_parse_code_trailing_newlines_expression():
+    code = """\
+x = 5
+y = 10
+x + y
+
+
+"""
+    code, setup, code_type = parse_code(code)
+
+    assert code == "x + y"
+    # strip() removes leading/trailing empty lines
+    assert setup == "x = 5\ny = 10"
+    assert code_type == "expression"
+
+
+def test_parse_code_trailing_whitespace_expression():
+    code, setup, code_type = parse_code("2 + 2    \n\n   ")
+
+    # strip() removes trailing whitespace
+    assert code == "2 + 2"
+    assert setup == ""
+    assert code_type == "expression"
+
+
+def test_parse_code_leading_newlines_expression():
+    code, setup, code_type = parse_code("\n\n\n2 + 2")
+
+    # Single expressions are returned as-is, not stripped
+    assert code == "\n\n\n2 + 2"
+    assert setup == ""
+    assert code_type == "expression"
+
+
+def test_parse_code_multiline_trailing_newlines():
+    code, setup, code_type = parse_code("x = 5\nx + 10\n\n")
+
+    assert code == "x + 10"
+    assert setup == "x = 5"
+    assert code_type == "expression"
+
+
+def test_parse_code_empty_list():
+    code, setup, code_type = parse_code("[]")
+
+    assert code == "[]"
+    assert setup == ""
+    assert code_type == "expression"
+
+
+def test_filter_existing_imports_star_import():
+    result = filter_existing_imports("from os import *", {"os": True})
+
+    assert result == "from os import *"
+
+
+def test_filter_existing_imports_relative():
+    result = filter_existing_imports(
+        "from ..models import User\nfrom ...core import Base", {}
+    )
+
+    assert result == "from ..models import User\nfrom ...core import Base"
+
+
+def test_filter_existing_imports_invalid_raises_error():
+    with pytest.raises(ValueError, match="Input must contain only import statements"):
+        filter_existing_imports("x = 5", {})