Finir
8bfd50de08
- ponytail_find_reinvented_stdlib: detect reimpl. stdlib patterns - ponytail_find_dead_code: find unused files/symbols - ponytail_find_single_impl_interface: YAGNI abstract class detector - ponytail_find_ponytail_comments: extract ponytail: debt ledger - ponytail_analyze_bloat: file bloat analysis
526 lines
19 KiB
Python
526 lines
19 KiB
Python
"""Ponytail MCP Server — help agents simplify code."""
|
|
import ast
|
|
import os
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from mcp.server import Server, stdio_server
|
|
from mcp.types import TextContent, Tool
|
|
|
|
# --- helpers ---
|
|
|
|
PYTHON_KEYWORDS = frozenset(
|
|
"False None True and as assert async await break class continue "
|
|
"def del elif else except finally for from global if import in "
|
|
"is lambda nonlocal not or pass raise return try while with yield".split()
|
|
)
|
|
|
|
|
|
def walk_python_files(root: str) -> list[Path]:
|
|
"""Recursively find all .py files under root."""
|
|
found = []
|
|
for p in Path(root).rglob("*.py"):
|
|
try:
|
|
rel = p.relative_to(root)
|
|
except ValueError:
|
|
rel = p
|
|
found.append(rel)
|
|
return sorted(found)
|
|
|
|
|
|
def read_lines(path: str) -> list[str]:
|
|
try:
|
|
with open(path) as f:
|
|
return f.readlines()
|
|
except OSError:
|
|
return []
|
|
|
|
|
|
# --- tool 1: find reinvented stdlib ---
|
|
|
|
# patterns: function names / class names that reimplement stdlib
|
|
REINVENTION_PATTERNS = {
|
|
"flatten": re.compile(r"def\s+(flatten|flatten_list)\s*\("),
|
|
"average": re.compile(r"def\s+(average|avg|mean)\s*\("),
|
|
"url_parse": re.compile(
|
|
r"(def\s+(parse_url|url_parse|parse_uri|url_parser)\s*\(|"
|
|
r"class\s+(URLParser|URIParser|UrlParser))"
|
|
),
|
|
"deep_merge": re.compile(r"def\s+(deep_merge|merge_dicts)\s*\("),
|
|
"chunk": re.compile(r"def\s+(chunk|chunked|chunks)\s*\("),
|
|
"debounce": re.compile(r"def\s+(debounce|throttle)\s*\("),
|
|
"slugify": re.compile(r"def\s+(slugify|to_slug)\s*\("),
|
|
"unique": re.compile(r"def\s+(unique|deduplicate|distinct)\s*\("),
|
|
"flatten_json": re.compile(r"def\s+(flatten_json|flatten_dict)\s*\("),
|
|
"retry": re.compile(r"def\s+(retry|retry_with|with_retry)\s*\("),
|
|
}
|
|
|
|
|
|
def tool_find_reinvented_stdlib(_params: dict[str, Any]) -> list[dict[str, Any]]:
|
|
"""Find reimplemented Python standard library functions."""
|
|
cwd = os.getcwd()
|
|
results: list[dict[str, Any]] = []
|
|
for pyfile in walk_python_files(cwd):
|
|
lines = read_lines(os.path.join(cwd, str(pyfile)))
|
|
for i, line in enumerate(lines, 1):
|
|
stripped = line.strip()
|
|
if stripped.startswith("#") or stripped.startswith("'''") or stripped.startswith('"""'):
|
|
continue
|
|
for name, pattern in REINVENTION_PATTERNS.items():
|
|
if pattern.search(line):
|
|
results.append(
|
|
{"file": str(pyfile), "line": i, "code": stripped}
|
|
)
|
|
break
|
|
return results
|
|
|
|
|
|
# --- tool 2: find dead code ---
|
|
|
|
def _collect_imports(root: str, pyfiles: list[Path]) -> dict[str, str]:
|
|
"""Map module name -> relative path."""
|
|
imports: dict[str, str] = {}
|
|
for pyfile in pyfiles:
|
|
src = "\n".join(read_lines(os.path.join(root, str(pyfile))))
|
|
try:
|
|
tree = ast.parse(src, filename=str(pyfile))
|
|
except SyntaxError:
|
|
continue
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Import):
|
|
for alias in node.names:
|
|
mod = alias.asname or alias.name
|
|
imports[mod] = str(pyfile)
|
|
elif isinstance(node, ast.ImportFrom):
|
|
mod = node.module or ""
|
|
if node.level > 0:
|
|
parts = str(pyfile).split("/")[: -node.level]
|
|
mod = "/".join(parts) + ("/" + mod if mod else "")
|
|
imports[mod] = str(pyfile)
|
|
return imports
|
|
|
|
|
|
def tool_find_dead_code(_params: dict[str, Any]) -> list[dict[str, Any]]:
|
|
"""Find files never imported and symbols never referenced."""
|
|
cwd = os.getcwd()
|
|
pyfiles = walk_python_files(cwd)
|
|
if not pyfiles:
|
|
return []
|
|
|
|
# Build set of all imported module names
|
|
imports = _collect_imports(cwd, pyfiles)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
|
|
# Dead files: .py files not imported anywhere (excluding __init__)
|
|
for pyfile in pyfiles:
|
|
name = pyfile.stem
|
|
if name == "__init__":
|
|
continue
|
|
stem = pyfile.stem
|
|
if stem not in imports:
|
|
# Check if any file imports this module
|
|
found_import = False
|
|
for pyfile2 in pyfiles:
|
|
if pyfile2 == pyfile:
|
|
continue
|
|
src = "\n".join(read_lines(os.path.join(cwd, str(pyfile2))))
|
|
try:
|
|
tree = ast.parse(src, filename=str(pyfile2))
|
|
except SyntaxError:
|
|
continue
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Import):
|
|
for alias in node.names:
|
|
if alias.name == stem or alias.name.startswith(stem + "."):
|
|
found_import = True
|
|
break
|
|
elif isinstance(node, ast.ImportFrom):
|
|
if node.module and (node.module == stem or node.module.startswith(stem + ".")):
|
|
found_import = True
|
|
break
|
|
if not found_import:
|
|
results.append(
|
|
{
|
|
"file": str(pyfile),
|
|
"line": 1,
|
|
"symbol": stem,
|
|
"reason": "module never imported",
|
|
}
|
|
)
|
|
|
|
# Dead symbols: top-level functions/classes not referenced elsewhere
|
|
for pyfile in pyfiles:
|
|
src = "\n".join(read_lines(os.path.join(cwd, str(pyfile))))
|
|
try:
|
|
tree = ast.parse(src, filename=str(pyfile))
|
|
except SyntaxError:
|
|
continue
|
|
|
|
# Collect defined top-level names
|
|
defined: list[tuple[str, int]] = []
|
|
for node in ast.iter_child_nodes(tree):
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
defined.append((node.name, node.lineno))
|
|
elif isinstance(node, ast.ClassDef):
|
|
defined.append((node.name, node.lineno))
|
|
|
|
# Collect all Name references across the whole project
|
|
all_refs: set[str] = set()
|
|
for other in pyfiles:
|
|
other_src = "\n".join(read_lines(os.path.join(cwd, str(other))))
|
|
try:
|
|
other_tree = ast.parse(other_src, filename=str(other))
|
|
except SyntaxError:
|
|
continue
|
|
for n in ast.walk(other_tree):
|
|
if isinstance(n, ast.Name):
|
|
all_refs.add(n.id)
|
|
elif isinstance(n, ast.Attribute):
|
|
if isinstance(n.value, ast.Name):
|
|
all_refs.add(n.value.id)
|
|
|
|
# Also check self-references within same file (imports count as refs)
|
|
for imp_name in imports:
|
|
base = imp_name.split(".")[-1]
|
|
all_refs.add(base)
|
|
|
|
for sym_name, lineno in defined:
|
|
if sym_name.startswith("_"):
|
|
continue
|
|
if sym_name in PYTHON_KEYWORDS:
|
|
continue
|
|
# Check if referenced outside its own file
|
|
refs_outside = False
|
|
for other in pyfiles:
|
|
if other == pyfile:
|
|
continue
|
|
other_src = "\n".join(read_lines(os.path.join(cwd, str(other))))
|
|
try:
|
|
other_tree = ast.parse(other_src, filename=str(other))
|
|
except SyntaxError:
|
|
continue
|
|
for n in ast.walk(other_tree):
|
|
if isinstance(n, ast.Name) and n.id == sym_name:
|
|
refs_outside = True
|
|
break
|
|
if isinstance(n, ast.Attribute) and isinstance(n.value, ast.Name) and n.value.id == sym_name:
|
|
refs_outside = True
|
|
break
|
|
if refs_outside:
|
|
break
|
|
|
|
if not refs_outside:
|
|
results.append(
|
|
{
|
|
"file": str(pyfile),
|
|
"line": lineno,
|
|
"symbol": sym_name,
|
|
"reason": "symbol never referenced",
|
|
}
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
# --- tool 3: find single-impl interfaces ---
|
|
|
|
def tool_find_single_impl_interface(_params: dict[str, Any]) -> list[dict[str, Any]]:
|
|
"""Find abstract classes with exactly one concrete implementation."""
|
|
cwd = os.getcwd()
|
|
pyfiles = walk_python_files(cwd)
|
|
if not pyfiles:
|
|
return []
|
|
|
|
# Collect all abstract base classes
|
|
abstract_classes: dict[str, tuple[str, int]] = {}
|
|
concrete_implementations: dict[str, list[tuple[str, int]]] = {}
|
|
|
|
for pyfile in pyfiles:
|
|
src = "\n".join(read_lines(os.path.join(cwd, str(pyfile))))
|
|
try:
|
|
tree = ast.parse(src, filename=str(pyfile))
|
|
except SyntaxError:
|
|
continue
|
|
|
|
for node in ast.iter_child_nodes(tree):
|
|
if not isinstance(node, ast.ClassDef):
|
|
continue
|
|
|
|
# Check if it's an abstract class
|
|
is_abstract = False
|
|
for base in node.bases:
|
|
base_name = ""
|
|
if isinstance(base, ast.Name):
|
|
base_name = base.id
|
|
elif isinstance(base, ast.Attribute):
|
|
base_name = base.attr
|
|
|
|
if base_name in ("ABC", "ABCMeta"):
|
|
is_abstract = True
|
|
if isinstance(base, ast.Attribute):
|
|
full = f"{ast.dump(base)}"
|
|
if "ABCMeta" in full or "ABC" in full:
|
|
is_abstract = True
|
|
|
|
# Check for abstract methods (decorated with @abstractmethod)
|
|
has_abstract_methods = False
|
|
for child in ast.iter_child_nodes(node):
|
|
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
for dec in child.decorator_list:
|
|
dec_name = ""
|
|
if isinstance(dec, ast.Name):
|
|
dec_name = dec.id
|
|
elif isinstance(dec, ast.Attribute):
|
|
dec_name = dec.attr
|
|
if "abstractmethod" in dec_name.lower():
|
|
has_abstract_methods = True
|
|
break
|
|
|
|
class_name = node.name
|
|
if is_abstract or has_abstract_methods:
|
|
abstract_classes[class_name] = (str(pyfile), node.lineno)
|
|
concrete_implementations[class_name] = []
|
|
|
|
# Check if this class inherits from a known abstract class
|
|
for base in node.bases:
|
|
base_name = ""
|
|
if isinstance(base, ast.Name):
|
|
base_name = base.id
|
|
elif isinstance(base, ast.Attribute):
|
|
base_name = base.attr
|
|
if base_name in abstract_classes:
|
|
concrete_implementations[base_name].append(
|
|
(str(pyfile), node.lineno)
|
|
)
|
|
|
|
# Also check for ABCMeta metaclass usage
|
|
for pyfile in pyfiles:
|
|
src = "\n".join(read_lines(os.path.join(cwd, str(pyfile))))
|
|
try:
|
|
tree = ast.parse(src, filename=str(pyfile))
|
|
except SyntaxError:
|
|
continue
|
|
for node in ast.iter_child_nodes(tree):
|
|
if isinstance(node, ast.ClassDef):
|
|
for kw in node.keywords:
|
|
if kw.arg == "metaclass":
|
|
if isinstance(kw.value, ast.Name) and kw.value.id == "ABCMeta":
|
|
cn = node.name
|
|
if cn not in abstract_classes:
|
|
abstract_classes[cn] = (str(pyfile), node.lineno)
|
|
concrete_implementations[cn] = []
|
|
|
|
results: list[dict[str, Any]] = []
|
|
for cls_name, (file_path, line_no) in abstract_classes.items():
|
|
impls = concrete_implementations.get(cls_name, [])
|
|
if len(impls) == 1:
|
|
impl_file, impl_line = impls[0]
|
|
results.append(
|
|
{
|
|
"abstract": cls_name,
|
|
"implementation": impls[0][0].rsplit("/", 1)[-1].replace(".py", ""),
|
|
"file": impl_file,
|
|
"line": impl_line,
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
# --- tool 4: find ponytail comments ---
|
|
|
|
def tool_find_ponytail_comments(_params: dict[str, Any]) -> list[dict[str, Any]]:
|
|
"""Extract all //ponytail: or #ponytail: comments."""
|
|
cwd = os.getcwd()
|
|
results: list[dict[str, Any]] = []
|
|
ponytail_re = re.compile(r"(?://|#)ponytail:\s*(.*)")
|
|
|
|
for pyfile in walk_python_files(cwd):
|
|
lines = read_lines(os.path.join(cwd, str(pyfile)))
|
|
for i, line in enumerate(lines, 1):
|
|
# Skip string literals (simple check)
|
|
stripped = line.strip()
|
|
if stripped.startswith("#!") or stripped.startswith('"""') or stripped.startswith("'''"):
|
|
continue
|
|
match = ponytail_re.search(line)
|
|
if match:
|
|
results.append(
|
|
{
|
|
"file": str(pyfile),
|
|
"line": i,
|
|
"comment": match.group(1).strip(),
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
# --- tool 5: analyze bloat ---
|
|
|
|
def _nesting_depth(node: ast.AST) -> int:
|
|
"""Compute max nesting depth of control structures."""
|
|
max_depth = 0
|
|
for child in ast.iter_child_nodes(node):
|
|
if isinstance(child, (ast.For, ast.AsyncFor, ast.While, ast.If, ast.With,
|
|
ast.AsyncWith, ast.Try, ast.ExceptHandler)):
|
|
child_depth = _nesting_depth(child) + 1
|
|
max_depth = max(max_depth, child_depth)
|
|
else:
|
|
child_depth = _nesting_depth(child)
|
|
max_depth = max(max_depth, child_depth)
|
|
return max_depth
|
|
|
|
|
|
def _func_lengths(tree: ast.Module) -> list[int]:
|
|
"""Get lengths of all top-level and nested functions."""
|
|
lengths: list[int] = []
|
|
|
|
def walk(node: ast.AST):
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
end = max(
|
|
(getattr(n, "end_lineno", n.lineno) for n in ast.walk(node)),
|
|
default=node.lineno,
|
|
)
|
|
lengths.append(end - node.lineno + 1)
|
|
for child in ast.iter_child_nodes(node):
|
|
walk(child)
|
|
|
|
walk(tree)
|
|
return lengths
|
|
|
|
|
|
def tool_analyze_bloat(params: dict[str, Any]) -> dict[str, Any]:
|
|
"""Analyze a file for bloat: lines, function lengths, nesting, imports."""
|
|
filepath = params.get("path", "")
|
|
if not filepath:
|
|
return {"error": "path required"}
|
|
|
|
if not os.path.exists(filepath):
|
|
return {"error": f"file not found: {filepath}"}
|
|
|
|
lines = read_lines(filepath)
|
|
total_lines = len(lines)
|
|
if total_lines == 0:
|
|
return {
|
|
"file": filepath,
|
|
"total_lines": 0,
|
|
"avg_function_length": 0,
|
|
"max_nesting_depth": 0,
|
|
"import_count": 0,
|
|
}
|
|
|
|
src = "\n".join(lines)
|
|
try:
|
|
tree = ast.parse(src, filename=filepath)
|
|
except SyntaxError:
|
|
return {
|
|
"file": filepath,
|
|
"total_lines": total_lines,
|
|
"avg_function_length": 0,
|
|
"max_nesting_depth": 0,
|
|
"import_count": 0,
|
|
"parse_error": True,
|
|
}
|
|
|
|
func_lengths = _func_lengths(tree)
|
|
avg_func_len = (
|
|
sum(func_lengths) / len(func_lengths) if func_lengths else 0
|
|
)
|
|
|
|
max_nesting = _nesting_depth(tree)
|
|
|
|
import_count = 0
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, (ast.Import, ast.ImportFrom)):
|
|
import_count += 1
|
|
|
|
return {
|
|
"file": filepath,
|
|
"total_lines": total_lines,
|
|
"avg_function_length": round(avg_func_len, 1),
|
|
"max_nesting_depth": max_nesting,
|
|
"import_count": import_count,
|
|
}
|
|
|
|
|
|
# --- server setup ---
|
|
|
|
async def run_server() -> None:
|
|
"""Run the ponytail MCP server."""
|
|
server = Server("ponytail")
|
|
|
|
@server.list_tools()
|
|
async def list_tools() -> list[Tool]:
|
|
return [
|
|
Tool(
|
|
name="ponytail_find_reinvented_stdlib",
|
|
description="Find reimplemented Python standard library functions (flatten, average, URL parser, etc.)",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {},
|
|
},
|
|
),
|
|
Tool(
|
|
name="ponytail_find_dead_code",
|
|
description="Find files never imported and symbols never referenced",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {},
|
|
},
|
|
),
|
|
Tool(
|
|
name="ponytail_find_single_impl_interface",
|
|
description="Find abstract classes with exactly one concrete implementation",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {},
|
|
},
|
|
),
|
|
Tool(
|
|
name="ponytail_find_ponytail_comments",
|
|
description="Extract all //ponytail: or #ponytail: comments from the workspace",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {},
|
|
},
|
|
),
|
|
Tool(
|
|
name="ponytail_analyze_bloat",
|
|
description="Analyze a file for bloat metrics: total lines, avg function length, max nesting depth, import count",
|
|
inputSchema={
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {
|
|
"type": "string",
|
|
"description": "Path to the Python file to analyze",
|
|
}
|
|
},
|
|
"required": ["path"],
|
|
},
|
|
),
|
|
]
|
|
|
|
@server.call_tool()
|
|
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
|
|
handlers = {
|
|
"ponytail_find_reinvented_stdlib": tool_find_reinvented_stdlib,
|
|
"ponytail_find_dead_code": tool_find_dead_code,
|
|
"ponytail_find_single_impl_interface": tool_find_single_impl_interface,
|
|
"ponytail_find_ponytail_comments": tool_find_ponytail_comments,
|
|
"ponytail_analyze_bloat": tool_analyze_bloat,
|
|
}
|
|
handler = handlers.get(name)
|
|
if not handler:
|
|
return [TextContent(type="text", text=f"Unknown tool: {name}")]
|
|
result = handler(arguments or {})
|
|
return [TextContent(type="text", text=__import__("json").dumps(result, separators=(",", ":")))]
|
|
|
|
async with stdio_server() as (read_stream, write_stream):
|
|
await server.run(read_stream, write_stream)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
asyncio.run(run_server())
|