athena-code 0.0.14__py3-none-any.whl

athena/mcp_server.py

@@ -0,0 +1,215 @@
+"""MCP server that exposes Athena Code Knowledge tools to Claude Code.
+
+This server wraps the `ack` CLI tool, providing structured access to code
+navigation capabilities through the Model Context Protocol.
+"""
+
+import json
+import subprocess
+from typing import Any
+
+from mcp.server import Server
+from mcp.server.stdio import stdio_server
+from mcp.types import Tool, TextContent
+
+
+# Initialize MCP server
+app = Server("ack")
+
+
+@app.list_tools()
+async def list_tools() -> list[Tool]:
+    """Declare available tools for Claude Code."""
+    return [
+        Tool(
+            name="ack_locate",
+            description=(
+                "Find the location of a Python entity (function, class, or method). "
+                "Returns file path and line range. Currently supports Python files only. "
+                "Use this to locate code before reading files - knowing the exact line "
+                "range allows targeted code extraction with tools like sed."
+            ),
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "entity": {
+                        "type": "string",
+                        "description": "Name of the entity to locate (e.g., 'validateSession', 'UserModel')",
+                    }
+                },
+                "required": ["entity"],
+            },
+        ),
+        Tool(
+            name="ack_info",
+            description=(
+                "Get detailed information about a code entity including signature, "
+                "parameters, return type, docstring, and dependencies. Supports functions, "
+                "classes, methods, modules, and packages. Returns structured JSON with all "
+                "available metadata about the entity."
+            ),
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": (
+                            "Path to entity in format 'file_path:entity_name' for functions/classes/methods, "
+                            "'file_path' for module-level info, or 'directory_path' for package info. "
+                            "Examples: 'src/auth.py:validate_token', 'src/auth.py', 'src/models'"
+                        ),
+                    }
+                },
+                "required": ["location"],
+            },
+        ),
+    ]
+
+
+@app.call_tool()
+async def call_tool(name: str, arguments: Any) -> list[TextContent]:
+    """Handle tool calls by routing to appropriate CLI commands."""
+    if name == "ack_locate":
+        return await _handle_locate(arguments["entity"])
+    elif name == "ack_info":
+        return await _handle_info(arguments["location"])
+
+    raise ValueError(f"Unknown tool: {name}")
+
+
+async def _handle_locate(entity: str) -> list[TextContent]:
+    """Handle ack_locate tool calls.
+
+    Args:
+        entity: Name of the entity to locate
+
+    Returns:
+        List containing a single TextContent with JSON results
+    """
+    try:
+        # Call the CLI tool
+        result = subprocess.run(
+            ["ack", "locate", entity],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
+        # Parse JSON output from CLI
+        locations = json.loads(result.stdout)
+
+        if not locations:
+            return [
+                TextContent(
+                    type="text",
+                    text=f"No entities found with name '{entity}'",
+                )
+            ]
+
+        # Format results for Claude Code
+        formatted_results = []
+        for loc in locations:
+            kind = loc["kind"]
+            path = loc["path"]
+            start = loc["extent"]["start"]
+            end = loc["extent"]["end"]
+            formatted_results.append(
+                f"{kind} '{entity}' found in {path} (lines {start}-{end})"
+            )
+
+        return [
+            TextContent(
+                type="text",
+                text="\n".join(formatted_results),
+            )
+        ]
+
+    except subprocess.CalledProcessError as e:
+        error_msg = e.stderr.strip() if e.stderr else str(e)
+        return [
+            TextContent(
+                type="text",
+                text=f"Error running ack locate: {error_msg}",
+            )
+        ]
+    except json.JSONDecodeError as e:
+        return [
+            TextContent(
+                type="text",
+                text=f"Error parsing ack output: {e}",
+            )
+        ]
+    except Exception as e:
+        return [
+            TextContent(
+                type="text",
+                text=f"Unexpected error: {e}",
+            )
+        ]
+
+
+async def _handle_info(location: str) -> list[TextContent]:
+    """Handle ack_info tool calls.
+
+    Args:
+        location: Path to entity in format "file_path:entity_name",
+                 "file_path" for module-level info,
+                 or "directory_path" for package info
+
+    Returns:
+        List containing a single TextContent with JSON results
+    """
+    try:
+        # Call the CLI tool
+        result = subprocess.run(
+            ["ack", "info", location],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
+        # Parse JSON output from CLI
+        entity_info = json.loads(result.stdout)
+
+        # Return formatted JSON
+        return [
+            TextContent(
+                type="text",
+                text=json.dumps(entity_info, indent=2),
+            )
+        ]
+
+    except subprocess.CalledProcessError as e:
+        error_msg = e.stderr.strip() if e.stderr else str(e)
+        return [
+            TextContent(
+                type="text",
+                text=f"Error running ack info: {error_msg}",
+            )
+        ]
+    except json.JSONDecodeError as e:
+        return [
+            TextContent(
+                type="text",
+                text=f"Error parsing ack output: {e}",
+            )
+        ]
+    except Exception as e:
+        return [
+            TextContent(
+                type="text",
+                text=f"Unexpected error: {e}",
+            )
+        ]
+
+
+async def main():
+    """Run the MCP server."""
+    async with stdio_server() as (read_stream, write_stream):
+        await app.run(read_stream, write_stream, app.create_initialization_options())
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

athena/models.py

@@ -0,0 +1,90 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Location:
+    """Represents a line range in a source file (0-indexed, inclusive)."""
+    start: int
+    end: int
+
+
+@dataclass
+class Entity:
+    """Represents a code entity (function, class, or method) found in a file."""
+    kind: str
+    path: str
+    extent: Location
+    name: str = ""  # Entity name (for filtering, not included in JSON output)
+
+
+@dataclass
+class Parameter:
+    """Represents a function/method parameter."""
+    name: str
+    type: str | None = None  # None if no type hint
+    default: str | None = None  # None if no default value
+
+
+@dataclass
+class Signature:
+    """Represents a function/method signature."""
+    name: str
+    args: list[Parameter]
+    return_type: str | None = None  # None if no return annotation
+
+
+@dataclass
+class FunctionInfo:
+    """Information about a function."""
+    path: str
+    extent: Location
+    sig: Signature
+    summary: str | None = None
+
+
+@dataclass
+class ClassInfo:
+    """Information about a class."""
+    path: str
+    extent: Location
+    methods: list[str]  # Formatted method signatures
+    summary: str | None = None
+
+
+@dataclass
+class MethodInfo:
+    """Information about a method."""
+    name: str  # Qualified name: "ClassName.method_name"
+    path: str
+    extent: Location
+    sig: Signature
+    summary: str | None = None
+
+
+@dataclass
+class ModuleInfo:
+    """Information about a module."""
+    path: str
+    extent: Location
+    summary: str | None = None
+
+
+@dataclass
+class PackageInfo:
+    """Information about a package (directory with __init__.py)."""
+    path: str
+    summary: str | None = None
+
+
+@dataclass
+class EntityStatus:
+    """Status information for an entity's hash synchronization state."""
+    kind: str
+    path: str
+    extent: str  # Format: "start-end" or empty for packages/modules
+    recorded_hash: str | None  # Hash from docstring, None if no hash
+    calculated_hash: str  # Hash computed from AST
+
+
+# Union type for entity info
+EntityInfo = FunctionInfo | ClassInfo | MethodInfo | ModuleInfo | PackageInfo

athena/parsers/__init__.py

@@ -0,0 +1,22 @@
+"""Parser module for extracting entities from source code"""
+from pathlib import Path
+
+from athena.parsers.base import BaseParser
+from athena.parsers.python_parser import PythonParser
+
+
+def get_parser_for_file(file_path: Path) -> BaseParser | None:
+    """Get the appropriate parser for a file based on its extension.
+
+    Args:
+        file_path: Path to the source file
+
+    Returns:
+        Parser instance if the file type is supported, None otherwise
+    """
+    extension = file_path.suffix.lower()
+
+    if extension == ".py":
+        return PythonParser()
+
+    return None

athena/parsers/base.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+
+from athena.models import ClassInfo, Entity, FunctionInfo, MethodInfo, ModuleInfo
+
+
+class BaseParser(ABC):
+    """Abstract base class for language-specific entity parsers."""
+
+    @abstractmethod
+    def extract_entities(self, source_code: str, file_path: str) -> list[Entity]:
+        """Extract all entities (functions, classes, methods) from source code.
+
+        Args:
+            source_code: The source code to parse
+            file_path: Relative path to the file (for Entity.path)
+
+        Returns:
+            List of Entity objects found in the source code
+        """
+        pass
+
+    @abstractmethod
+    def extract_entity_info(
+        self,
+        source_code: str,
+        file_path: str,
+        entity_name: str | None = None
+    ) -> FunctionInfo | ClassInfo | MethodInfo | ModuleInfo | None:
+        """Extract detailed information about a specific entity.
+
+        Args:
+            source_code: The source code to parse
+            file_path: Relative path to the file (for EntityInfo.path)
+            entity_name: Name of entity to find, or None for module-level
+
+        Returns:
+            EntityInfo object, or None if not found
+        """
+        pass