mne-docs-mcp 1.0.0

package/python/ast_extractor.py

@@ -0,0 +1,422 @@
+import ast
+from typing import Optional, List, Dict, Any
+
+def extract_docstring_params(docstring: str) -> List[Dict[str, str]]:
+    """Parse numpy-style docstring to extract parameter info."""
+    params = []
+    if not docstring:
+        return params
+        
+    lines = docstring.split('\n')
+    in_params = False
+    current_param = None
+    
+    for line in lines:
+        line = line.strip()
+        if line == 'Parameters':
+            in_params = True
+            continue
+        if line.startswith('----------') or line == 'Returns':
+            if in_params:
+                break
+        
+        if in_params and line:
+            # Check for parameter definition "name : type"
+            if ' : ' in line:
+                parts = line.split(' : ')
+                name = parts[0].strip()
+                type_info = parts[1].strip()
+                current_param = {'name': name, 'type': type_info, 'description': ''}
+                params.append(current_param)
+            elif current_param:
+                # Append description
+                if current_param['description']:
+                    current_param['description'] += ' ' + line
+                else:
+                    current_param['description'] = line
+                    
+    return params
+
+
+def extract_docstring_examples(docstring: str) -> List[str]:
+    """Extract code examples from numpy-style docstring."""
+    examples = []
+    if not docstring:
+        return examples
+    
+    lines = docstring.split('\n')
+    in_examples = False
+    skip_separator = False
+    current_example = []
+    
+    for line in lines:
+        stripped = line.strip()
+        
+        # Check for Examples section
+        if stripped == 'Examples':
+            in_examples = True
+            skip_separator = True  # Skip the next line (-------)
+            continue
+        
+        # Skip the separator line after section header
+        if skip_separator:
+            skip_separator = False
+            if stripped.startswith('---'):
+                continue
+        
+        # Check for end of Examples section (next section header)
+        if in_examples and stripped in ('See Also', 'Notes', 'References', 'Attributes', 'Methods', 'Returns', 'Parameters'):
+            break
+            
+        if in_examples:
+            # Detect start of code block with >>>
+            if stripped.startswith('>>>'):
+                if current_example:
+                    examples.append('\n'.join(current_example))
+                current_example = [stripped[4:] if stripped.startswith('>>> ') else stripped[3:]]
+            elif stripped.startswith('...'):
+                # Continuation line
+                if current_example:
+                    current_example.append(stripped[4:] if stripped.startswith('... ') else stripped[3:])
+            elif current_example and stripped and not stripped.startswith('---'):
+                # Output line - include it as comment
+                current_example.append('# Output: ' + stripped)
+    
+    # Don't forget the last example
+    if current_example:
+        examples.append('\n'.join(current_example))
+    
+    return examples[:5]  # Limit to 5 examples
+
+def get_decorators(node: ast.AST) -> List[str]:
+    decorators = []
+    if hasattr(node, 'decorator_list'):
+        for decorator in node.decorator_list:
+            if isinstance(decorator, ast.Name):
+                decorators.append(decorator.id)
+            elif isinstance(decorator, ast.Call):
+                if isinstance(decorator.func, ast.Name):
+                    decorators.append(decorator.func.id)
+                elif isinstance(decorator.func, ast.Attribute):
+                    decorators.append(decorator.func.attr)
+    return decorators
+
+def get_signature(node: ast.FunctionDef) -> str:
+    args = []
+    # Handle positional args
+    for arg in node.args.args:
+        arg_str = arg.arg
+        if arg.annotation:
+            if isinstance(arg.annotation, ast.Name):
+                arg_str += f": {arg.annotation.id}"
+            elif isinstance(arg.annotation, ast.Constant):
+                arg_str += f": {arg.annotation.value}"
+        args.append(arg_str)
+        
+    # Handle defaults (naive approach, matching from end)
+    if node.args.defaults:
+        for i, default in enumerate(reversed(node.args.defaults)):
+            if isinstance(default, ast.Constant):
+                args[-(i+1)] += f"={default.value}"
+                
+    return f"({', '.join(args)})"
+
+class SymbolVisitor(ast.NodeVisitor):
+    def __init__(self, module_name: str = "", file_url: str = ""):
+        self.symbols = []
+        self.module_name = module_name
+        self.file_url = file_url
+        self.current_class = None
+
+    def visit_ClassDef(self, node):
+        qualified_name = f"{self.module_name}.{node.name}" if self.module_name else node.name
+        
+        symbol = {
+            'name': node.name,
+            'qualifiedName': qualified_name,
+            'kind': 'class',
+            'signature': f"class {node.name}",
+            'docstring': ast.get_docstring(node) or "",
+            'lineno': node.lineno,
+            'endLineno': node.end_lineno,
+            'decorators': get_decorators(node),
+            'baseClasses': [b.id for b in node.bases if isinstance(b, ast.Name)],
+            'parameters': extract_docstring_params(ast.get_docstring(node) or "")
+        }
+        self.symbols.append(symbol)
+        
+        old_class = self.current_class
+        self.current_class = qualified_name
+        self.generic_visit(node)
+        self.current_class = old_class
+
+    def visit_FunctionDef(self, node):
+        is_method = self.current_class is not None
+        name = node.name
+        qualified_name = f"{self.current_class}.{name}" if is_method else (f"{self.module_name}.{name}" if self.module_name else name)
+        
+        symbol = {
+            'name': name,
+            'qualifiedName': qualified_name,
+            'kind': 'method' if is_method else 'function',
+            'signature': f"def {name}{get_signature(node)}",
+            'docstring': ast.get_docstring(node) or "",
+            'lineno': node.lineno,
+            'endLineno': node.end_lineno,
+            'decorators': get_decorators(node),
+            'parameters': extract_docstring_params(ast.get_docstring(node) or "")
+        }
+        self.symbols.append(symbol)
+        self.generic_visit(node)
+
+def parse_file(source: str, module_name: str = "", file_url: str = "") -> List[Dict[str, Any]]:
+    try:
+        tree = ast.parse(source)
+        visitor = SymbolVisitor(module_name, file_url)
+        visitor.visit(tree)
+        return visitor.symbols
+    except Exception:
+        return []
+
+def extract_symbol(source: str, symbol_name: str, module_name: str = "", file_url: str = "") -> Optional[Dict[str, Any]]:
+    symbols = parse_file(source, module_name, file_url)
+    result = None
+    
+    # Try exact match first
+    for s in symbols:
+        if s['name'] == symbol_name or s['qualifiedName'] == symbol_name:
+            result = s
+            break
+            
+    # Try partial match for methods (e.g. "Epochs.plot" matching "plot" inside "Epochs")
+    if result is None and '.' in symbol_name:
+        parts = symbol_name.split('.')
+        short_name = parts[-1]
+        for s in symbols:
+            if s['name'] == short_name and (s['qualifiedName'].endswith(symbol_name) or symbol_name.endswith(s['qualifiedName'])):
+                result = s
+                break
+    
+    if result:
+        # Extract examples from docstring
+        examples = extract_docstring_examples(result.get('docstring', ''))
+        if examples:
+            result['examples'] = examples
+    
+        # If we found a class, also include its methods with full details
+        if result['kind'] == 'class':
+            class_qualified_name = result['qualifiedName']
+            methods = []
+            for s in symbols:
+                if s['kind'] == 'method' and s['qualifiedName'].startswith(class_qualified_name + '.'):
+                    # Include __init__ and public methods
+                    if not s['name'].startswith('_') or s['name'] == '__init__':
+                        # Get a meaningful docstring summary (first line or first sentence)
+                        docstring = s.get('docstring', '')
+                        if docstring:
+                            # Get first non-empty line as summary
+                            lines = [l.strip() for l in docstring.split('\n') if l.strip()]
+                            summary = lines[0] if lines else ''
+                            # Truncate if too long
+                            if len(summary) > 200:
+                                summary = summary[:197] + '...'
+                        else:
+                            summary = ''
+                        
+                        methods.append({
+                            'name': s['name'],
+                            'signature': s['signature'],
+                            'docstring': summary,
+                            'lineno': s.get('lineno'),
+                            'parameters': s.get('parameters', [])[:5],  # First 5 params
+                        })
+            
+            # Sort methods: __init__ first, then alphabetically
+            methods.sort(key=lambda m: (0 if m['name'] == '__init__' else 1, m['name']))
+            result['methods'] = methods[:30]  # Limit to 30 methods
+                
+    return result
+
+def find_symbols(source: str, query: str, module_name: str = "", max_results: int = 20) -> List[Dict[str, Any]]:
+    symbols = parse_file(source, module_name)
+    query = query.lower()
+    matches = []
+    
+    for s in symbols:
+        if query in s['name'].lower() or query in s['qualifiedName'].lower():
+            matches.append(s)
+            
+    return matches[:max_results]
+
+
+class CrossReferenceVisitor(ast.NodeVisitor):
+    """Extract function calls and attribute accesses from source code."""
+    
+    def __init__(self, module_name: str = ""):
+        self.module_name = module_name
+        self.imports = {}  # alias -> full module path
+        self.from_imports = {}  # name -> (module, original_name)
+        self.calls = []  # List of function/method calls
+        self.current_function = None
+        self.current_class = None
+    
+    def visit_Import(self, node):
+        for alias in node.names:
+            name = alias.asname or alias.name
+            self.imports[name] = alias.name
+        self.generic_visit(node)
+    
+    def visit_ImportFrom(self, node):
+        module = node.module or ''
+        for alias in node.names:
+            name = alias.asname or alias.name
+            self.from_imports[name] = (module, alias.name)
+        self.generic_visit(node)
+    
+    def visit_ClassDef(self, node):
+        old_class = self.current_class
+        self.current_class = node.name
+        self.generic_visit(node)
+        self.current_class = old_class
+    
+    def visit_FunctionDef(self, node):
+        old_func = self.current_function
+        if self.current_class:
+            self.current_function = f"{self.current_class}.{node.name}"
+        else:
+            self.current_function = node.name
+        self.generic_visit(node)
+        self.current_function = old_func
+    
+    visit_AsyncFunctionDef = visit_FunctionDef
+    
+    def visit_Call(self, node):
+        call_info = self._resolve_call(node)
+        if call_info:
+            call_info['caller'] = self.current_function
+            call_info['lineno'] = node.lineno
+            self.calls.append(call_info)
+        self.generic_visit(node)
+    
+    def _resolve_call(self, node) -> Optional[Dict[str, Any]]:
+        """Resolve a Call node to a callee name."""
+        func = node.func
+        
+        if isinstance(func, ast.Name):
+            # Direct function call: func()
+            name = func.id
+            if name in self.from_imports:
+                module, orig_name = self.from_imports[name]
+                return {'callee': f"{module}.{orig_name}", 'kind': 'function'}
+            return {'callee': name, 'kind': 'function'}
+        
+        elif isinstance(func, ast.Attribute):
+            # Method/attribute call: obj.method()
+            parts = []
+            current = func
+            while isinstance(current, ast.Attribute):
+                parts.append(current.attr)
+                current = current.value
+            
+            if isinstance(current, ast.Name):
+                parts.append(current.id)
+                parts.reverse()
+                
+                # Check if first part is an import alias
+                first = parts[0]
+                if first in self.imports:
+                    parts[0] = self.imports[first]
+                elif first in self.from_imports:
+                    module, orig_name = self.from_imports[first]
+                    parts[0] = f"{module}.{orig_name}"
+                
+                return {'callee': '.'.join(parts), 'kind': 'method'}
+        
+        return None
+
+
+def extract_cross_references(source: str, symbol_name: str, module_name: str = "") -> Dict[str, Any]:
+    """
+    Extract cross-references for a symbol: what it calls (callees) and what calls it (callers).
+    
+    symbol_name can be:
+      - "function_name" for module-level functions
+      - "ClassName.method_name" for class methods
+      - "ClassName" for all methods in a class
+    """
+    try:
+        tree = ast.parse(source)
+    except SyntaxError:
+        return {'symbol': symbol_name, 'callees': [], 'callers': [], 'error': 'Syntax error in source'}
+    
+    visitor = CrossReferenceVisitor(module_name)
+    visitor.visit(tree)
+    
+    # Parse symbol_name to handle Class.method format
+    class_context = None
+    method_name = symbol_name
+    if '.' in symbol_name:
+        parts = symbol_name.split('.')
+        # Check if first part looks like a class (PascalCase)
+        if parts[0] and parts[0][0].isupper():
+            class_context = parts[0]
+            method_name = parts[1] if len(parts) > 1 else None
+    
+    # Find callees (what the symbol calls)
+    callees = []
+    # Find callers (what calls the symbol)
+    callers = []
+    
+    for call in visitor.calls:
+        caller = call.get('caller')
+        callee = call.get('callee', '')
+        
+        # Check if this call is FROM our symbol (symbol is the caller)
+        caller_matches = False
+        if caller:
+            if class_context and method_name:
+                # Looking for ClassName.method_name
+                caller_matches = caller == f"{class_context}.{method_name}"
+            elif class_context and not method_name:
+                # Looking for any method in ClassName
+                caller_matches = caller.startswith(f"{class_context}.")
+            else:
+                # Looking for a function name
+                caller_matches = (caller == symbol_name or 
+                                caller.endswith(f".{symbol_name}"))
+        
+        if caller_matches:
+            callees.append({
+                'name': callee,
+                'kind': call.get('kind'),
+                'lineno': call.get('lineno')
+            })
+        
+        # Check if this call is TO our symbol (symbol is the callee)
+        callee_matches = (callee == symbol_name or 
+                        callee.endswith(f".{symbol_name}") or
+                        (method_name and callee.endswith(f".{method_name}")))
+        
+        if callee_matches:
+            callers.append({
+                'name': caller or '<module>',
+                'kind': 'function' if caller else 'module',
+                'lineno': call.get('lineno')
+            })
+    
+    # Deduplicate callees by name
+    seen_callees = set()
+    unique_callees = []
+    for c in callees:
+        if c['name'] not in seen_callees:
+            seen_callees.add(c['name'])
+            unique_callees.append(c)
+    
+    return {
+        'symbol': symbol_name,
+        'module': module_name,
+        'callees': unique_callees[:50],  # Limit results
+        'callers': callers[:50],
+        'totalCallees': len(unique_callees),
+        'totalCallers': len(callers)
+    }

package/python/parser.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""
+Python AST parser for MNE documentation extraction.
+Communicates with TypeScript via stdin/stdout JSON.
+"""
+
+import ast
+import json
+import sys
+import traceback
+from typing import Optional, List, Dict, Any
+from ast_extractor import extract_symbol, find_symbols, parse_file, extract_cross_references
+
+def main():
+    """Main entry point - read command from stdin, write result to stdout."""
+    # Ensure stdout is using utf-8
+    sys.stdout.reconfigure(encoding='utf-8')
+    sys.stdin.reconfigure(encoding='utf-8')
+
+    try:
+        # Read all input from stdin
+        input_data = sys.stdin.read()
+        if not input_data:
+            return
+
+        command = json.loads(input_data)
+        action = command.get('action')
+        
+        result = None
+        
+        if action == 'parse_file':
+            result = parse_file(
+                command['source'],
+                command.get('moduleName', ''),
+                command.get('fileUrl', '')
+            )
+        elif action == 'extract_symbol':
+            result = extract_symbol(
+                command['source'],
+                command['symbolName'],
+                command.get('moduleName', ''),
+                command.get('fileUrl', '')
+            )
+        elif action == 'find_symbols':
+            result = find_symbols(
+                command['source'],
+                command['query'],
+                command.get('moduleName', ''),
+                command.get('maxResults', 20)
+            )
+        elif action == 'extract_cross_references':
+            result = extract_cross_references(
+                command['source'],
+                command['symbolName'],
+                command.get('moduleName', '')
+            )
+        else:
+            result = {'error': f'Unknown action: {action}'}
+        
+        json.dump({'success': True, 'data': result}, sys.stdout)
+    
+    except Exception as e:
+        # Capture full traceback for debugging
+        tb = traceback.format_exc()
+        json.dump({'success': False, 'error': str(e), 'traceback': tb}, sys.stdout)
+
+if __name__ == '__main__':
+    main()

package/python/requirements.txt

@@ -0,0 +1,2 @@
+# No external dependencies needed for AST parsing
+# Standard library only: ast, json, sys, typing