feat(mcp): Extract return type hints and ImGui scopes in code outline
This commit is contained in:
+35
-4
@@ -53,7 +53,11 @@ class CodeOutliner:
|
||||
return doc.splitlines()[0]
|
||||
return None
|
||||
|
||||
count = [0]
|
||||
def walk(node: ast.AST, indent: int = 0) -> None:
|
||||
count[0] += 1
|
||||
if count[0] > 100000:
|
||||
raise Exception("Infinite loop detected! " + str(type(node)))
|
||||
"""
|
||||
[C: src/summarize.py:_summarise_python]
|
||||
"""
|
||||
@@ -70,15 +74,42 @@ class CodeOutliner:
|
||||
start_line = node.lineno
|
||||
end_line = getattr(node, "end_lineno", start_line)
|
||||
prefix = "[Async Func]" if isinstance(node, ast.AsyncFunctionDef) else "[Func]"
|
||||
# Check if it's a method
|
||||
# We can check the indent or the parent, but in AST walk we know if we are inside a ClassDef
|
||||
# Let's use a simpler heuristic for the outline: if indent > 0, it's likely a method.
|
||||
if indent > 0:
|
||||
prefix = "[Method]"
|
||||
output.append(f"{' ' * indent}{prefix} {node.name} (Lines {start_line}-{end_line})")
|
||||
returns = ""
|
||||
if getattr(node, "returns", None):
|
||||
try:
|
||||
returns = f" -> {ast.unparse(node.returns)}"
|
||||
except Exception:
|
||||
pass
|
||||
output.append(f"{' ' * indent}{prefix} {node.name}{returns} (Lines {start_line}-{end_line})")
|
||||
doc = get_docstring(node)
|
||||
if doc:
|
||||
output.append(f"{' ' * (indent + 1)}\"\"\"{doc}\"\"\"")
|
||||
for item in node.body:
|
||||
walk(item, indent + 1)
|
||||
elif isinstance(node, ast.With):
|
||||
is_imgui = False
|
||||
try:
|
||||
for item in node.items:
|
||||
ctx_str = ast.unparse(item.context_expr)
|
||||
if "imscope." in ctx_str or "imgui." in ctx_str:
|
||||
start_line = node.lineno
|
||||
end_line = getattr(node, "end_lineno", start_line)
|
||||
output.append(f"{' ' * indent}[ImGui Scope] {ctx_str} (Lines {start_line}-{end_line})")
|
||||
is_imgui = True
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
for item in node.body:
|
||||
walk(item, indent + 1 if is_imgui else indent)
|
||||
else:
|
||||
for block_attr in ("body", "orelse", "handlers", "finalbody"):
|
||||
block = getattr(node, block_attr, [])
|
||||
if isinstance(block, list):
|
||||
for item in block:
|
||||
walk(item, indent)
|
||||
|
||||
for node in tree.body:
|
||||
walk(node)
|
||||
return "\n".join(output)
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
import pytest
|
||||
import ast
|
||||
from pathlib import Path
|
||||
from src.outline_tool import CodeOutliner
|
||||
|
||||
def test_code_outliner_type_hints():
|
||||
code = """
|
||||
class Test:
|
||||
def my_method(self, x: int) -> str:
|
||||
return "hello"
|
||||
|
||||
def my_func(a: bool) -> None:
|
||||
pass
|
||||
"""
|
||||
outliner = CodeOutliner()
|
||||
result = outliner.outline(code)
|
||||
assert "[Method] my_method -> str (Lines 3-4)" in result
|
||||
assert "[Func] my_func -> None (Lines 6-7)" in result
|
||||
|
||||
def test_code_outliner_imgui_scopes():
|
||||
code = """
|
||||
def render_gui() -> None:
|
||||
with imscope.window("My Window"):
|
||||
imgui.text("Hello")
|
||||
with imscope.child("My Child"):
|
||||
pass
|
||||
"""
|
||||
outliner = CodeOutliner()
|
||||
result = outliner.outline(code)
|
||||
assert "[ImGui Scope] imscope.window('My Window') (Lines 3-6)" in result
|
||||
assert "[ImGui Scope] imscope.child('My Child') (Lines 5-6)" in result
|
||||
|
||||
def test_code_outliner_nested_ifs():
|
||||
code = """
|
||||
def render_gui() -> None:
|
||||
if True:
|
||||
with imscope.window("My Window"):
|
||||
pass
|
||||
"""
|
||||
outliner = CodeOutliner()
|
||||
result = outliner.outline(code)
|
||||
assert "[ImGui Scope] imscope.window('My Window') (Lines 4-5)" in result
|
||||
Reference in New Issue
Block a user