lean-lsp-mcp 0.14.1__py3-none-any.whl → 0.16.0__py3-none-any.whl

lean_lsp_mcp/outline_utils.py

@@ -3,43 +3,52 @@ from typing import Dict, List, Optional, Tuple
 from leanclient import LeanLSPClient
 from leanclient.utils import DocumentContentChange
 
+from lean_lsp_mcp.models import FileOutline, OutlineEntry
+
 
 METHOD_KIND = {6, "method"}
 KIND_TAGS = {"namespace": "Ns"}
 
 
-def _get_info_trees(client: LeanLSPClient, path: str, symbols: List[Dict]) -> Dict[str, str]:
+def _get_info_trees(
+    client: LeanLSPClient, path: str, symbols: List[Dict]
+) -> Dict[str, str]:
     """Insert #info_trees commands, collect diagnostics, then revert changes."""
     if not symbols:
         return {}
 
     symbol_by_line = {}
     changes = []
-    for i, sym in enumerate(sorted(symbols, key=lambda s: s['range']['start']['line'])):
-        line = sym['range']['start']['line'] + i
-        symbol_by_line[line] = sym['name']
+    for i, sym in enumerate(sorted(symbols, key=lambda s: s["range"]["start"]["line"])):
+        line = sym["range"]["start"]["line"] + i
+        symbol_by_line[line] = sym["name"]
         changes.append(DocumentContentChange("#info_trees in\n", [line, 0], [line, 0]))
 
     client.update_file(path, changes)
     diagnostics = client.get_diagnostics(path)
 
     info_trees = {
-        symbol_by_line[diag['range']['start']['line']]: diag['message']
+        symbol_by_line[diag["range"]["start"]["line"]]: diag["message"]
         for diag in diagnostics
-        if diag['severity'] == 3 and diag['range']['start']['line'] in symbol_by_line
+        if diag["severity"] == 3 and diag["range"]["start"]["line"] in symbol_by_line
     }
 
     # Revert in reverse order
-    client.update_file(path, [
-        DocumentContentChange("", [line, 0], [line + 1, 0])
-        for line in sorted(symbol_by_line.keys(), reverse=True)
-    ])
+    client.update_file(
+        path,
+        [
+            DocumentContentChange("", [line, 0], [line + 1, 0])
+            for line in sorted(symbol_by_line.keys(), reverse=True)
+        ],
+    )
     return info_trees
 
 
 def _extract_type(info: str, name: str) -> Optional[str]:
     """Extract type signature from info tree message."""
-    if m := re.search(rf'  • \[Term\] {re.escape(name)} \(isBinder := true\) : ([^@]+) @', info):
+    if m := re.search(
+        rf"  • \[Term\] {re.escape(name)} \(isBinder := true\) : ([^@]+) @", info
+    ):
         return m.group(1).strip()
     return None
 
@@ -47,14 +56,16 @@ def _extract_type(info: str, name: str) -> Optional[str]:
 def _extract_fields(info: str, name: str) -> List[Tuple[str, str]]:
     """Extract structure/class fields from info tree message."""
     fields = []
-    for pattern in [rf'{re.escape(name)}\.(\w+)', rf'@{re.escape(name)}\.(\w+)']:
-        for m in re.finditer(rf'  • \[Term\] {pattern} \(isBinder := true\) : (.+?) @', info):
+    for pattern in [rf"{re.escape(name)}\.(\w+)", rf"@{re.escape(name)}\.(\w+)"]:
+        for m in re.finditer(
+            rf"  • \[Term\] {pattern} \(isBinder := true\) : (.+?) @", info
+        ):
             field_name, full_type = m.groups()
             # Clean up the type signature
-            if ']' in full_type:
-                field_type = full_type[full_type.rfind(']')+1:].lstrip('→ ').strip()
-            elif ' → ' in full_type:
-                field_type = full_type.split(' → ')[-1].strip()
+            if "]" in full_type:
+                field_type = full_type[full_type.rfind("]") + 1 :].lstrip("→ ").strip()
+            elif " → " in full_type:
+                field_type = full_type.split(" → ")[-1].strip()
             else:
                 field_type = full_type.strip()
             fields.append((field_name, field_type))
@@ -68,51 +79,61 @@ def _extract_declarations(content: str, start: int, end: int) -> List[Dict]:
 
     while i < min(end, len(lines)):
         line = lines[i].strip()
-        for keyword in ['theorem', 'lemma', 'def']:
+        for keyword in ["theorem", "lemma", "def"]:
             if line.startswith(f"{keyword} "):
-                name = line[len(keyword):].strip().split()[0]
-                if name and not name.startswith('_'):
+                name = line[len(keyword) :].strip().split()[0]
+                if name and not name.startswith("_"):
                     # Collect until :=
                     decl_lines = [line]
                     j = i + 1
-                    while j < min(end, len(lines)) and ':=' not in ' '.join(decl_lines):
-                        if (next_line := lines[j].strip()) and not next_line.startswith('--'):
+                    while j < min(end, len(lines)) and ":=" not in " ".join(decl_lines):
+                        if (next_line := lines[j].strip()) and not next_line.startswith(
+                            "--"
+                        ):
                             decl_lines.append(next_line)
                         j += 1
 
                     # Extract signature (everything before :=, minus keyword and name)
-                    full_decl = ' '.join(decl_lines)
+                    full_decl = " ".join(decl_lines)
                     type_sig = None
-                    if ':=' in full_decl:
-                        sig_part = full_decl.split(':=', 1)[0].strip()[len(keyword):].strip()
+                    if ":=" in full_decl:
+                        sig_part = (
+                            full_decl.split(":=", 1)[0].strip()[len(keyword) :].strip()
+                        )
                         if sig_part.startswith(name):
-                            type_sig = sig_part[len(name):].strip()
-
-                    decls.append({
-                        'name': name,
-                        'kind': 'method',
-                        'range': {'start': {'line': i, 'character': 0}, 
-                                 'end': {'line': i, 'character': len(lines[i])}},
-                        '_keyword': keyword,
-                        '_type': type_sig
-                    })
+                            type_sig = sig_part[len(name) :].strip()
+
+                    decls.append(
+                        {
+                            "name": name,
+                            "kind": "method",
+                            "range": {
+                                "start": {"line": i, "character": 0},
+                                "end": {"line": i, "character": len(lines[i])},
+                            },
+                            "_keyword": keyword,
+                            "_type": type_sig,
+                        }
+                    )
                 break
         i += 1
     return decls
 
 
-def _flatten_symbols(symbols: List[Dict], indent: int = 0, content: str = "") -> List[Tuple[Dict, int]]:
+def _flatten_symbols(
+    symbols: List[Dict], indent: int = 0, content: str = ""
+) -> List[Tuple[Dict, int]]:
     """Recursively flatten symbol hierarchy, extracting declarations from namespaces."""
     result = []
     for sym in symbols:
         result.append((sym, indent))
-        children = sym.get('children', [])
+        children = sym.get("children", [])
 
         # Extract theorem/lemma/def from namespace bodies
-        if content and sym.get('kind') == 'namespace':
-            ns_range = sym['range']
-            ns_start = ns_range['start']['line']
-            ns_end = ns_range['end']['line']
+        if content and sym.get("kind") == "namespace":
+            ns_range = sym["range"]
+            ns_start = ns_range["start"]["line"]
+            ns_end = ns_range["end"]["line"]
             children = children + _extract_declarations(content, ns_start, ns_end)
 
         if children:
@@ -120,32 +141,36 @@ def _flatten_symbols(symbols: List[Dict], indent: int = 0, content: str = "") ->
     return result
 
 
-def _detect_tag(name: str, kind: str, type_sig: str, has_fields: bool, keyword: Optional[str]) -> str:
+def _detect_tag(
+    name: str, kind: str, type_sig: str, has_fields: bool, keyword: Optional[str]
+) -> str:
     """Determine the appropriate tag for a symbol."""
     if has_fields:
-        return "Class" if '→' in type_sig else "Struct"
+        return "Class" if "→" in type_sig else "Struct"
     if name == "example":
         return "Ex"
-    if keyword in {'theorem', 'lemma'}:
+    if keyword in {"theorem", "lemma"}:
         return "Thm"
-    if type_sig and any(marker in type_sig for marker in ['∀', '=']):
+    if type_sig and any(marker in type_sig for marker in ["∀", "="]):
         return "Thm"
-    if type_sig and '→' in type_sig.replace(' → ', '', 1):  # More than one arrow
+    if type_sig and "→" in type_sig.replace(" → ", "", 1):  # More than one arrow
         return "Thm"
     return KIND_TAGS.get(kind, "Def")
 
 
 def _format_symbol(sym: Dict, type_sigs: Dict, fields_map: Dict, indent: int) -> str:
     """Format a single symbol with its type signature and fields."""
-    name = sym['name']
-    type_sig = sym.get('_type') or type_sigs.get(name, "")
+    name = sym["name"]
+    type_sig = sym.get("_type") or type_sigs.get(name, "")
     fields = fields_map.get(name, [])
 
-    tag = _detect_tag(name, sym.get('kind', ''), type_sig, bool(fields), sym.get('_keyword'))
+    tag = _detect_tag(
+        name, sym.get("kind", ""), type_sig, bool(fields), sym.get("_keyword")
+    )
     prefix = "\t" * indent
 
-    start = sym['range']['start']['line'] + 1
-    end = sym['range']['end']['line'] + 1
+    start = sym["range"]["start"]["line"] + 1
+    end = sym["range"]["end"]["line"] + 1
     line_info = f"L{start}" if start == end else f"L{start}-{end}"
 
     result = f"{prefix}[{tag}: {line_info}] {name}"
@@ -158,14 +183,108 @@ def _format_symbol(sym: Dict, type_sigs: Dict, fields_map: Dict, indent: int) ->
     return result + "\n"
 
 
+def _build_outline_entry(
+    sym: Dict, type_sigs: Dict, fields_map: Dict, indent: int
+) -> Optional[OutlineEntry]:
+    """Build a structured outline entry for a symbol."""
+    name = sym["name"]
+    type_sig = sym.get("_type") or type_sigs.get(name, "")
+    fields = fields_map.get(name, [])
+
+    tag = _detect_tag(
+        name, sym.get("kind", ""), type_sig, bool(fields), sym.get("_keyword")
+    )
+    start = sym["range"]["start"]["line"] + 1
+    end = sym["range"]["end"]["line"] + 1
+
+    # Add fields as children for structs/classes
+    children = [
+        OutlineEntry(
+            name=fname,
+            kind="field",
+            start_line=start,
+            end_line=start,
+            type_signature=ftype,
+            children=[],
+        )
+        for fname, ftype in fields
+    ]
+
+    return OutlineEntry(
+        name=name,
+        kind=tag,
+        start_line=start,
+        end_line=end,
+        type_signature=type_sig if type_sig else None,
+        children=children,
+    )
+
+
+def generate_outline_data(client: LeanLSPClient, path: str) -> FileOutline:
+    """Generate structured outline data for a Lean file."""
+    client.open_file(path)
+    content = client.get_file_content(path)
+
+    # Extract imports
+    imports = [
+        line.strip()[7:]
+        for line in content.splitlines()
+        if line.strip().startswith("import ")
+    ]
+
+    symbols = client.get_document_symbols(path)
+    if not symbols and not imports:
+        return FileOutline(imports=[], declarations=[])
+
+    # Flatten symbol tree and extract namespace declarations
+    all_symbols = _flatten_symbols(symbols, content=content)
+
+    # Get info trees only for LSP symbols (not extracted declarations)
+    lsp_methods = [
+        s
+        for s, _ in all_symbols
+        if s.get("kind") in METHOD_KIND and "_keyword" not in s
+    ]
+    info_trees = _get_info_trees(client, path, lsp_methods)
+
+    # Extract type signatures and fields from info trees
+    type_sigs = {
+        name: sig
+        for name, info in info_trees.items()
+        if (sig := _extract_type(info, name))
+    }
+    fields_map = {
+        name: fields
+        for name, info in info_trees.items()
+        if (fields := _extract_fields(info, name))
+    }
+
+    # Build declarations list
+    declarations = []
+    for sym, indent in all_symbols:
+        if (
+            sym.get("kind") in METHOD_KIND
+            or sym.get("_keyword")
+            or sym.get("kind") == "namespace"
+        ):
+            entry = _build_outline_entry(sym, type_sigs, fields_map, indent)
+            if entry:
+                declarations.append(entry)
+
+    return FileOutline(imports=imports, declarations=declarations)
+
+
 def generate_outline(client: LeanLSPClient, path: str) -> str:
     """Generate a concise outline of a Lean file showing structure and signatures."""
     client.open_file(path)
     content = client.get_file_content(path)
 
     # Extract imports
-    imports = [line.strip()[7:] for line in content.splitlines() 
-               if line.strip().startswith("import ")]
+    imports = [
+        line.strip()[7:]
+        for line in content.splitlines()
+        if line.strip().startswith("import ")
+    ]
 
     symbols = client.get_document_symbols(path)
     if not symbols and not imports:
@@ -175,14 +294,24 @@ def generate_outline(client: LeanLSPClient, path: str) -> str:
     all_symbols = _flatten_symbols(symbols, content=content)
 
     # Get info trees only for LSP symbols (not extracted declarations)
-    lsp_methods = [s for s, _ in all_symbols if s.get('kind') in METHOD_KIND and '_keyword' not in s]
+    lsp_methods = [
+        s
+        for s, _ in all_symbols
+        if s.get("kind") in METHOD_KIND and "_keyword" not in s
+    ]
     info_trees = _get_info_trees(client, path, lsp_methods)
 
     # Extract type signatures and fields from info trees
-    type_sigs = {name: sig for name, info in info_trees.items() 
-                 if (sig := _extract_type(info, name))}
-    fields_map = {name: fields for name, info in info_trees.items() 
-                  if (fields := _extract_fields(info, name))}
+    type_sigs = {
+        name: sig
+        for name, info in info_trees.items()
+        if (sig := _extract_type(info, name))
+    }
+    fields_map = {
+        name: fields
+        for name, info in info_trees.items()
+        if (fields := _extract_fields(info, name))
+    }
 
     # Build output
     parts = []
@@ -193,7 +322,9 @@ def generate_outline(client: LeanLSPClient, path: str) -> str:
         declarations = [
             _format_symbol(sym, type_sigs, fields_map, indent)
             for sym, indent in all_symbols
-            if sym.get('kind') in METHOD_KIND or sym.get('_keyword') or sym.get('kind') == 'namespace'
+            if sym.get("kind") in METHOD_KIND
+            or sym.get("_keyword")
+            or sym.get("kind") == "namespace"
         ]
         parts.append("## Declarations\n" + "".join(declarations).rstrip())