feat: ponytail MCP server — 5 code simplification tools
- 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
This commit is contained in:
Finir
@@ -0,0 +1,525 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user