87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
import re
|
|
import ast
|
|
import sqlglot
|
|
from typing import Dict, List, Any
|
|
from pathlib import Path
|
|
|
|
def extract_sql_schema(sql_code: str) -> Dict[str, List[str]]:
|
|
"""
|
|
Extracts table -> columns mapping from SQL code using sqlglot.
|
|
Returns a canonical schema.
|
|
"""
|
|
schema: Dict[str, List[str]] = {}
|
|
try:
|
|
statements = sqlglot.parse(sql_code, read="postgres")
|
|
except Exception:
|
|
return schema
|
|
|
|
for stmt in statements:
|
|
if stmt.key and stmt.key.upper() == "CREATE":
|
|
for table in stmt.find_all(sqlglot.exp.Create):
|
|
try:
|
|
tname = table.this.this
|
|
cols = []
|
|
for coldef in table.find_all(sqlglot.exp.ColumnDef):
|
|
cname = getattr(coldef.this, "name", None)
|
|
if cname:
|
|
cols.append(cname)
|
|
if tname and cols:
|
|
schema[tname] = cols
|
|
except Exception:
|
|
continue
|
|
return schema
|
|
|
|
|
|
def extract_python_structure(code: str) -> Dict[str, List[str]]:
|
|
"""Return module structure: functions and classes."""
|
|
try:
|
|
tree = ast.parse(code)
|
|
except SyntaxError:
|
|
return {}
|
|
|
|
funcs = [n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)]
|
|
classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
|
|
return {"functions": funcs, "classes": classes}
|
|
|
|
|
|
def extract_go_structure(code: str) -> Dict[str, List[str]]:
|
|
"""Simple regex-based Go structure detection (lightweight)."""
|
|
funcs = re.findall(r"func\s+([A-Z]\w+)", code)
|
|
structs = re.findall(r"type\s+(\w+)\s+struct", code)
|
|
return {"functions": funcs, "structs": structs}
|
|
|
|
|
|
def extract_rust_structure(code: str) -> Dict[str, List[str]]:
|
|
"""Heuristic Rust index."""
|
|
structs = re.findall(r"struct\s+(\w+)", code)
|
|
traits = re.findall(r"trait\s+(\w+)", code)
|
|
funcs = re.findall(r"fn\s+(\w+)", code)
|
|
return {"functions": funcs, "structs": structs, "traits": traits}
|
|
|
|
|
|
def extract_svelte_structure(code: str) -> Dict[str, List[str]]:
|
|
"""Minimal Svelte export/prop finder."""
|
|
exports = re.findall(r"export\s+let\s+(\w+)", code)
|
|
funcs = re.findall(r"function\s+(\w+)", code)
|
|
return {"props": exports, "functions": funcs}
|
|
|
|
|
|
def build_quick_index(language: str, code: str, filepath: Path) -> Dict[str, Any]:
|
|
"""Dispatch to the appropriate structure extractor."""
|
|
data = {}
|
|
if language == "sql":
|
|
data = extract_sql_schema(code)
|
|
elif language == "python":
|
|
data = extract_python_structure(code)
|
|
elif language == "go":
|
|
data = extract_go_structure(code)
|
|
elif language == "rust":
|
|
data = extract_rust_structure(code)
|
|
elif language == "svelte":
|
|
data = extract_svelte_structure(code)
|
|
|
|
return {
|
|
"text": f"Quick index for {filepath.name}:\n{data}",
|
|
"metadata": {"language": language, "file": str(filepath), "type": "index"}
|
|
}
|