heimdall-code 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- heimdall/SKILL.md +54 -0
- heimdall/__init__.py +13 -0
- heimdall/__main__.py +8 -0
- heimdall/cli.py +200 -0
- heimdall/graph.py +286 -0
- heimdall/index.py +435 -0
- heimdall/install.py +151 -0
- heimdall/mcp_server.py +157 -0
- heimdall/parse.py +409 -0
- heimdall/registry.py +51 -0
- heimdall_code-0.1.0.dist-info/METADATA +293 -0
- heimdall_code-0.1.0.dist-info/RECORD +15 -0
- heimdall_code-0.1.0.dist-info/WHEEL +4 -0
- heimdall_code-0.1.0.dist-info/entry_points.txt +3 -0
- heimdall_code-0.1.0.dist-info/licenses/LICENSE +202 -0
heimdall/index.py
ADDED
|
@@ -0,0 +1,435 @@
|
|
|
1
|
+
# Copyright 2026 The Heimdall Authors
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
"""The indexing pipeline: discover -> parse -> resolve -> store.
|
|
4
|
+
|
|
5
|
+
Four passes, run once per repository. Everything a query needs is precomputed here
|
|
6
|
+
so the read side stays a single SQL statement.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import os
|
|
12
|
+
import subprocess
|
|
13
|
+
import time
|
|
14
|
+
from concurrent.futures import ProcessPoolExecutor
|
|
15
|
+
from concurrent.futures.process import BrokenProcessPool
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
|
|
18
|
+
from .graph import Graph
|
|
19
|
+
from .parse import EXT_LANG, FileParse, parse_source
|
|
20
|
+
|
|
21
|
+
MAX_BYTES = 2_000_000
|
|
22
|
+
SKIP_DIRS = {
|
|
23
|
+
".git",
|
|
24
|
+
".hg",
|
|
25
|
+
".svn",
|
|
26
|
+
"node_modules",
|
|
27
|
+
"venv",
|
|
28
|
+
".venv",
|
|
29
|
+
"__pycache__",
|
|
30
|
+
".mypy_cache",
|
|
31
|
+
".ruff_cache",
|
|
32
|
+
".pytest_cache",
|
|
33
|
+
"dist",
|
|
34
|
+
"build",
|
|
35
|
+
"target",
|
|
36
|
+
".tox",
|
|
37
|
+
".idea",
|
|
38
|
+
".gradle",
|
|
39
|
+
"site-packages",
|
|
40
|
+
".next",
|
|
41
|
+
".terraform",
|
|
42
|
+
"vendor",
|
|
43
|
+
"coverage-report",
|
|
44
|
+
}
|
|
45
|
+
# package roots and index files that do not contribute a module name segment
|
|
46
|
+
_INDEX_STEMS = {"__init__", "index", "mod", "lib", "main"}
|
|
47
|
+
_SRC_ROOTS = {"src", "lib", "app", "pkg", "internal", "source"}
|
|
48
|
+
# receivers that name the enclosing class rather than another object
|
|
49
|
+
SELF_NAMES = {"self", "this", "cls", "me", "Self"}
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# ---------------------------------------------------------------- pass 1: discover
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def discover(root: Path) -> list[Path]:
|
|
56
|
+
"""Tracked + untracked-but-not-ignored files when git knows; a plain walk otherwise."""
|
|
57
|
+
try:
|
|
58
|
+
out = subprocess.run(
|
|
59
|
+
["git", "-C", str(root), "ls-files", "-z", "--cached", "--others", "--exclude-standard"],
|
|
60
|
+
capture_output=True,
|
|
61
|
+
timeout=120,
|
|
62
|
+
check=True,
|
|
63
|
+
).stdout
|
|
64
|
+
paths = [root / p.decode() for p in out.split(b"\0") if p]
|
|
65
|
+
# a superproject lists each submodule as one directory entry, so a repo-of-repos
|
|
66
|
+
# would index to nothing at all. --recurse-submodules only works with --cached,
|
|
67
|
+
# hence the second call rather than a flag on the first.
|
|
68
|
+
if any(p.is_dir() for p in paths):
|
|
69
|
+
out = subprocess.run(
|
|
70
|
+
["git", "-C", str(root), "ls-files", "-z", "--cached", "--recurse-submodules"],
|
|
71
|
+
capture_output=True,
|
|
72
|
+
timeout=120,
|
|
73
|
+
check=True,
|
|
74
|
+
).stdout
|
|
75
|
+
paths += [root / p.decode() for p in out.split(b"\0") if p]
|
|
76
|
+
except (subprocess.SubprocessError, OSError):
|
|
77
|
+
paths = []
|
|
78
|
+
for dirpath, dirnames, filenames in os.walk(root):
|
|
79
|
+
dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")]
|
|
80
|
+
paths.extend(Path(dirpath) / f for f in filenames)
|
|
81
|
+
|
|
82
|
+
keep = []
|
|
83
|
+
for p in paths:
|
|
84
|
+
if p.suffix not in EXT_LANG or any(part in SKIP_DIRS for part in p.parts):
|
|
85
|
+
continue
|
|
86
|
+
try:
|
|
87
|
+
if p.is_file() and p.stat().st_size <= MAX_BYTES:
|
|
88
|
+
keep.append(p)
|
|
89
|
+
except OSError:
|
|
90
|
+
continue
|
|
91
|
+
return sorted(set(keep))
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
# ---------------------------------------------------------------- pass 2: parse
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def _parse_one(args: tuple[str, str]) -> FileParse:
|
|
98
|
+
rel, abspath = args
|
|
99
|
+
lang = EXT_LANG[Path(rel).suffix]
|
|
100
|
+
try:
|
|
101
|
+
src = Path(abspath).read_bytes()
|
|
102
|
+
except OSError as exc:
|
|
103
|
+
fp = FileParse(path=rel, lang=lang)
|
|
104
|
+
fp.error = str(exc)
|
|
105
|
+
return fp
|
|
106
|
+
return parse_source(rel, lang, src)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def parse_all(root: Path, files: list[Path], jobs: int, log) -> list[FileParse]:
|
|
110
|
+
work = [(str(p.relative_to(root)), str(p)) for p in files]
|
|
111
|
+
if jobs == 1:
|
|
112
|
+
return [_parse_one(w) for w in work]
|
|
113
|
+
results: list[FileParse] = []
|
|
114
|
+
try:
|
|
115
|
+
with ProcessPoolExecutor(max_workers=jobs) as pool:
|
|
116
|
+
for i, fp in enumerate(pool.map(_parse_one, work, chunksize=16), 1):
|
|
117
|
+
results.append(fp)
|
|
118
|
+
if i % 200 == 0:
|
|
119
|
+
log(f" parsed {i}/{len(work)} files")
|
|
120
|
+
except BrokenProcessPool:
|
|
121
|
+
# a REPL, a notebook or a frozen binary: the workers cannot import __main__
|
|
122
|
+
log(" worker pool unavailable, parsing in this process")
|
|
123
|
+
return [_parse_one(w) for w in work]
|
|
124
|
+
return results
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
# ---------------------------------------------------------------- pass 3: resolve
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def _module_keys(rel: str) -> list[str]:
|
|
131
|
+
"""Every dotted suffix a file could be imported as, longest first."""
|
|
132
|
+
p = Path(rel)
|
|
133
|
+
dirs = list(p.parts[:-1])
|
|
134
|
+
# `main.py` is both `pkg.main` (python) and `pkg` (js/go entry point); register both,
|
|
135
|
+
# named form first so the longer, more specific import wins.
|
|
136
|
+
forms = [[*dirs, p.stem]] + ([dirs] if p.stem in _INDEX_STEMS else [])
|
|
137
|
+
keys: list[str] = []
|
|
138
|
+
for parts in forms:
|
|
139
|
+
parts = list(parts)
|
|
140
|
+
while parts and parts[0] in _SRC_ROOTS:
|
|
141
|
+
parts.pop(0)
|
|
142
|
+
keys += [".".join(parts[i:]) for i in range(len(parts))]
|
|
143
|
+
return keys
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
_IMPORT_EXTS = (".js", ".ts", ".jsx", ".tsx", ".py", ".h", ".hpp", ".mjs", ".cjs")
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def _strip_ext(mod: str) -> str:
|
|
150
|
+
for ext in _IMPORT_EXTS:
|
|
151
|
+
if mod.endswith(ext):
|
|
152
|
+
return mod[: -len(ext)]
|
|
153
|
+
return mod
|
|
154
|
+
|
|
155
|
+
|
|
156
|
+
def _normalize_import(mod: str) -> str:
|
|
157
|
+
mod = _strip_ext(mod.strip().strip("\"'<>"))
|
|
158
|
+
return mod.replace("::", ".").replace("/", ".").replace("\\", ".").strip(".")
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def _relative_target(from_path: str, mod: str) -> str | None:
|
|
162
|
+
"""`./store` and `from ..pkg.x import y` are relative to the importing file's
|
|
163
|
+
directory. Returns a repo-relative path stem, or None if the import is absolute."""
|
|
164
|
+
mod = _strip_ext(mod.strip().strip("\"'<>"))
|
|
165
|
+
ups = len(mod) - len(mod.lstrip("."))
|
|
166
|
+
if ups == 0:
|
|
167
|
+
return None
|
|
168
|
+
rest = mod[ups:].lstrip("/").replace(".", "/")
|
|
169
|
+
base = os.path.dirname(from_path)
|
|
170
|
+
for _ in range(ups - 1):
|
|
171
|
+
base = os.path.dirname(base)
|
|
172
|
+
return os.path.normpath(os.path.join(base, rest)).replace(os.sep, "/")
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def build(root: Path, db_path: Path, jobs: int = 0, log=lambda m: None) -> dict:
|
|
176
|
+
t0 = time.time()
|
|
177
|
+
# process_cpu_count respects the affinity mask, but only exists on 3.13+
|
|
178
|
+
count = getattr(os, "process_cpu_count", os.cpu_count)
|
|
179
|
+
jobs = jobs or min(32, count() or 4)
|
|
180
|
+
|
|
181
|
+
log("pass 1/4 discovering files")
|
|
182
|
+
files = discover(root)
|
|
183
|
+
log(f" {len(files)} source files")
|
|
184
|
+
if not files:
|
|
185
|
+
raise SystemExit(f"no supported source files under {root}")
|
|
186
|
+
|
|
187
|
+
log(f"pass 2/4 parsing ({jobs} workers)")
|
|
188
|
+
parsed = parse_all(root, files, jobs, log)
|
|
189
|
+
|
|
190
|
+
log("pass 3/4 building graph")
|
|
191
|
+
g = Graph.open(db_path)
|
|
192
|
+
# the index lives inside the repo it describes; keep it out of the repo's git status
|
|
193
|
+
(db_path.parent / ".gitignore").write_text("*\n")
|
|
194
|
+
g.db.executescript("DELETE FROM edge; DELETE FROM node; DELETE FROM meta;")
|
|
195
|
+
|
|
196
|
+
# --- folder + file nodes
|
|
197
|
+
nodes: list[dict] = []
|
|
198
|
+
folders: set[str] = set()
|
|
199
|
+
for fp in parsed:
|
|
200
|
+
d = Path(fp.path).parent
|
|
201
|
+
while str(d) not in (".", "/"):
|
|
202
|
+
folders.add(str(d))
|
|
203
|
+
d = d.parent
|
|
204
|
+
for f in sorted(folders):
|
|
205
|
+
nodes.append(
|
|
206
|
+
dict(kind="Folder", name=Path(f).name, qname=f, path=f, line=0, end_line=0, lang=None, signature="")
|
|
207
|
+
)
|
|
208
|
+
for fp in parsed:
|
|
209
|
+
nodes.append(
|
|
210
|
+
dict(
|
|
211
|
+
kind="File",
|
|
212
|
+
name=Path(fp.path).name,
|
|
213
|
+
qname=fp.path,
|
|
214
|
+
path=fp.path,
|
|
215
|
+
line=0,
|
|
216
|
+
end_line=0,
|
|
217
|
+
lang=fp.lang,
|
|
218
|
+
signature="",
|
|
219
|
+
)
|
|
220
|
+
)
|
|
221
|
+
for s in fp.symbols:
|
|
222
|
+
nodes.append(
|
|
223
|
+
dict(
|
|
224
|
+
kind=s.kind,
|
|
225
|
+
name=s.name,
|
|
226
|
+
qname=s.qname,
|
|
227
|
+
path=fp.path,
|
|
228
|
+
line=s.line,
|
|
229
|
+
end_line=s.end_line,
|
|
230
|
+
lang=fp.lang,
|
|
231
|
+
signature=s.signature,
|
|
232
|
+
)
|
|
233
|
+
)
|
|
234
|
+
g.add_nodes(nodes)
|
|
235
|
+
ids = g.ids_by_qname()
|
|
236
|
+
|
|
237
|
+
edges: list[tuple] = []
|
|
238
|
+
# --- containment
|
|
239
|
+
for f in folders:
|
|
240
|
+
parent = str(Path(f).parent)
|
|
241
|
+
if parent in ids:
|
|
242
|
+
edges.append((ids[parent], ids[f], "CONTAINS", 1.0, 0))
|
|
243
|
+
for fp in parsed:
|
|
244
|
+
parent = str(Path(fp.path).parent)
|
|
245
|
+
if parent in ids:
|
|
246
|
+
edges.append((ids[parent], ids[fp.path], "CONTAINS", 1.0, 0))
|
|
247
|
+
for s in fp.symbols:
|
|
248
|
+
src = ids.get(s.parent) if s.parent else ids.get(fp.path)
|
|
249
|
+
if src and s.qname in ids:
|
|
250
|
+
edges.append((src, ids[s.qname], "DEFINES", 1.0, s.line))
|
|
251
|
+
|
|
252
|
+
# --- imports: relative paths resolve against the importing file, absolute ones by
|
|
253
|
+
# longest dotted suffix
|
|
254
|
+
module_index: dict[str, int] = {}
|
|
255
|
+
path_index: dict[str, int] = {}
|
|
256
|
+
for fp in parsed:
|
|
257
|
+
fid = ids[fp.path]
|
|
258
|
+
for key in _module_keys(fp.path):
|
|
259
|
+
module_index.setdefault(key, fid)
|
|
260
|
+
stem = fp.path.rsplit(".", 1)[0]
|
|
261
|
+
path_index.setdefault(stem, fid)
|
|
262
|
+
if Path(fp.path).stem in _INDEX_STEMS:
|
|
263
|
+
path_index.setdefault(str(Path(fp.path).parent), fid)
|
|
264
|
+
imported: dict[str, set[str]] = {}
|
|
265
|
+
# per file: lines of imports that landed on a file in this repo, so an `as` rename
|
|
266
|
+
# from a third-party module is not mistaken for a local symbol
|
|
267
|
+
inside: dict[str, set[int]] = {}
|
|
268
|
+
# per file, per enclosing symbol (None = module level): the files it imports
|
|
269
|
+
scoped: dict[str, dict[str | None, set[str]]] = {}
|
|
270
|
+
# per file: the receiver name a module was imported under -> that module's path,
|
|
271
|
+
# so `utils.helper()` resolves to utils.py instead of every `helper` in the repo.
|
|
272
|
+
aliases: dict[str, dict[str, str]] = {}
|
|
273
|
+
path_by_id = {fid: fp.path for fp in parsed if (fid := ids[fp.path])}
|
|
274
|
+
for fp in parsed:
|
|
275
|
+
fid = ids[fp.path]
|
|
276
|
+
seen = imported.setdefault(fp.path, set())
|
|
277
|
+
here_lines = inside.setdefault(fp.path, set())
|
|
278
|
+
alias = aliases.setdefault(fp.path, {})
|
|
279
|
+
scopes = scoped.setdefault(fp.path, {})
|
|
280
|
+
for mod, line, scope in fp.imports:
|
|
281
|
+
rel = _relative_target(fp.path, mod)
|
|
282
|
+
if rel is not None:
|
|
283
|
+
target = path_index.get(rel)
|
|
284
|
+
if target is not None and target != fid:
|
|
285
|
+
edges.append((fid, target, "IMPORTS", 1.0, line))
|
|
286
|
+
seen.add(path_by_id[target])
|
|
287
|
+
here_lines.add(line)
|
|
288
|
+
scopes.setdefault(scope, set()).add(path_by_id[target])
|
|
289
|
+
alias[rel.rsplit("/", 1)[-1]] = path_by_id[target]
|
|
290
|
+
continue
|
|
291
|
+
key = _normalize_import(mod)
|
|
292
|
+
while key:
|
|
293
|
+
target = module_index.get(key)
|
|
294
|
+
if target is not None and target != fid:
|
|
295
|
+
edges.append((fid, target, "IMPORTS", 1.0, line))
|
|
296
|
+
seen.add(path_by_id[target])
|
|
297
|
+
here_lines.add(line)
|
|
298
|
+
scopes.setdefault(scope, set()).add(path_by_id[target])
|
|
299
|
+
alias[key.rsplit(".", 1)[-1]] = path_by_id[target]
|
|
300
|
+
break
|
|
301
|
+
key = key.split(".", 1)[1] if "." in key else ""
|
|
302
|
+
|
|
303
|
+
# --- name index for call/heritage resolution
|
|
304
|
+
by_name: dict[str, list[tuple[int, str, str, str, str]]] = {}
|
|
305
|
+
for fp in parsed:
|
|
306
|
+
for sym in fp.symbols:
|
|
307
|
+
if sym.kind in ("Function", "Method", "Class", "Interface") and sym.qname in ids:
|
|
308
|
+
by_name.setdefault(sym.name, []).append((ids[sym.qname], fp.path, sym.kind, sym.qname, fp.lang))
|
|
309
|
+
|
|
310
|
+
# --- heritage
|
|
311
|
+
for fp in parsed:
|
|
312
|
+
for cls_qname, base in fp.heritage:
|
|
313
|
+
src = ids.get(cls_qname)
|
|
314
|
+
if src is None:
|
|
315
|
+
continue
|
|
316
|
+
ranked = _rank(by_name.get(base, ()), fp.path, imported.get(fp.path, ()), lang=fp.lang)
|
|
317
|
+
for dst, _path, _kind, conf in ranked[:1]:
|
|
318
|
+
if dst != src:
|
|
319
|
+
edges.append((src, dst, "EXTENDS", conf, 0))
|
|
320
|
+
|
|
321
|
+
# --- calls
|
|
322
|
+
ambiguous = 0
|
|
323
|
+
for fp in parsed:
|
|
324
|
+
fid = ids[fp.path]
|
|
325
|
+
local = imported.get(fp.path, ())
|
|
326
|
+
alias = aliases.get(fp.path, {})
|
|
327
|
+
scopes = scoped.get(fp.path, {})
|
|
328
|
+
# only renames whose import resolved inside this repo: `from asyncio import run
|
|
329
|
+
# as arun` must not turn `arun()` into a call to some local `run`
|
|
330
|
+
here_lines = inside.get(fp.path, ())
|
|
331
|
+
renames = {loc: orig for loc, orig, ln in fp.renames if ln in here_lines}
|
|
332
|
+
for caller, callee, line, receiver in fp.calls:
|
|
333
|
+
src = ids.get(caller) or fid
|
|
334
|
+
# a caller with its own imports is resolved against those, not the file's:
|
|
335
|
+
# dispatchers that import a different `execute` per function are common
|
|
336
|
+
own = scopes.get(caller)
|
|
337
|
+
here = (own | scopes.get(None, set())) if own else local
|
|
338
|
+
# `from x import y as z` -- the call says `z`, the symbol is called `y`
|
|
339
|
+
if callee not in by_name and callee in renames:
|
|
340
|
+
callee = renames[callee]
|
|
341
|
+
# `self.foo()` inside a method can only mean a sibling member of that class.
|
|
342
|
+
# ponytail: own class only -- inherited members fall through to name ranking.
|
|
343
|
+
scope: tuple[str, ...] = ()
|
|
344
|
+
scope_path = None
|
|
345
|
+
if receiver in SELF_NAMES and caller and "." in caller:
|
|
346
|
+
scope = (caller.rsplit(".", 1)[0] + ".",)
|
|
347
|
+
elif receiver:
|
|
348
|
+
scope_path = alias.get(receiver) or alias.get(renames.get(receiver, ""))
|
|
349
|
+
if scope_path is None:
|
|
350
|
+
# `OrderTable.full_name()` -- the receiver names a class, so the
|
|
351
|
+
# target is that class's member and not every same-named method
|
|
352
|
+
scope = _class_scopes(by_name, renames.get(receiver, receiver), fp.path, here)
|
|
353
|
+
ranked = _rank(by_name.get(callee, ()), fp.path, here, scope, scope_path, fp.lang)
|
|
354
|
+
if not ranked:
|
|
355
|
+
continue
|
|
356
|
+
if len(ranked) > 1 and ranked[0][3] < 0.7:
|
|
357
|
+
ambiguous += 1
|
|
358
|
+
for dst, _path, _kind, conf in ranked[:5]:
|
|
359
|
+
edges.append((src, dst, "CALLS", conf, line)) # a self-edge is recursion
|
|
360
|
+
|
|
361
|
+
g.add_edges(edges)
|
|
362
|
+
log("pass 4/4 building search index")
|
|
363
|
+
g.reindex_fts()
|
|
364
|
+
g.set_meta(
|
|
365
|
+
root=str(root.resolve()),
|
|
366
|
+
files=len(files),
|
|
367
|
+
indexed_at=time.strftime("%Y-%m-%dT%H:%M:%S"),
|
|
368
|
+
parse_errors=sum(1 for fp in parsed if fp.error),
|
|
369
|
+
ambiguous_calls=ambiguous,
|
|
370
|
+
seconds=round(time.time() - t0, 2),
|
|
371
|
+
)
|
|
372
|
+
g.db.commit()
|
|
373
|
+
g.db.execute("PRAGMA optimize")
|
|
374
|
+
stats = g.stats()
|
|
375
|
+
g.close()
|
|
376
|
+
log(f"done in {stats['seconds']}s")
|
|
377
|
+
return stats
|
|
378
|
+
|
|
379
|
+
|
|
380
|
+
def _class_scopes(by_name, receiver: str, from_path: str, here) -> tuple[str, ...]:
|
|
381
|
+
"""qname prefixes of every class the receiver could name, nearest first.
|
|
382
|
+
|
|
383
|
+
Ambiguity in the *receiver* is not ambiguity in the *call*: Django has 236 classes
|
|
384
|
+
called `Model` and exactly one of them defines `check`, so `Model.check()` still has
|
|
385
|
+
one answer. Returning every candidate and letting the member decide is what makes
|
|
386
|
+
that work; returning None the moment the name is shared threw the receiver away.
|
|
387
|
+
"""
|
|
388
|
+
cands = [c for c in by_name.get(receiver, ()) if c[2] in ("Class", "Interface")]
|
|
389
|
+
ordered = sorted(cands, key=lambda c: (c[1] != from_path, c[1] not in here))
|
|
390
|
+
return tuple(c[3] + "." for c in ordered)
|
|
391
|
+
|
|
392
|
+
|
|
393
|
+
def _rank(
|
|
394
|
+
candidates,
|
|
395
|
+
from_path: str,
|
|
396
|
+
imported_paths,
|
|
397
|
+
scope: tuple[str, ...] = (),
|
|
398
|
+
scope_path: str | None = None,
|
|
399
|
+
lang: str | None = None,
|
|
400
|
+
) -> list[tuple]:
|
|
401
|
+
"""Score call targets: a known receiver beats same file beats imported file beats a
|
|
402
|
+
unique global name; anything else is an ambiguous name match and says so."""
|
|
403
|
+
cands = list(candidates)
|
|
404
|
+
if not cands:
|
|
405
|
+
return []
|
|
406
|
+
if lang and len({c[4] for c in cands}) > 1: # a Go `Store` is not a Python `Store`
|
|
407
|
+
cands = [c for c in cands if c[4] == lang] or cands
|
|
408
|
+
if scope: # receiver named a class: the target is a member of one of those classes
|
|
409
|
+
in_scope = [c for c in cands if c[3].startswith(scope)]
|
|
410
|
+
if in_scope:
|
|
411
|
+
in_scope.sort(key=lambda c: (c[1] != from_path, c[1] not in imported_paths))
|
|
412
|
+
# one class with that name defines it: certain. Several do: nearest, and say so.
|
|
413
|
+
return [(c[0], c[1], c[2], 0.99 if len(in_scope) == 1 else 0.9) for c in in_scope[:1]]
|
|
414
|
+
if scope_path: # receiver was an imported module: the target lives in that file
|
|
415
|
+
in_module = [c for c in cands if c[1] == scope_path]
|
|
416
|
+
if in_module:
|
|
417
|
+
return [(c[0], c[1], c[2], 0.9) for c in in_module[:1]]
|
|
418
|
+
imported_paths = set(imported_paths)
|
|
419
|
+
scored = []
|
|
420
|
+
for nid, path, kind, _qname, _lang in cands:
|
|
421
|
+
if path == from_path:
|
|
422
|
+
conf = 0.95
|
|
423
|
+
elif path in imported_paths:
|
|
424
|
+
conf = 0.8
|
|
425
|
+
elif len(cands) == 1:
|
|
426
|
+
conf = 0.7
|
|
427
|
+
else:
|
|
428
|
+
# ponytail: name-only fallback. Type inference on receivers would lift these;
|
|
429
|
+
# until then they are marked low-confidence and callers can filter them out.
|
|
430
|
+
conf = 0.3
|
|
431
|
+
scored.append((nid, path, kind, conf))
|
|
432
|
+
scored.sort(key=lambda x: -x[3])
|
|
433
|
+
# evidence beats a guess: once one candidate is in this file or an imported one,
|
|
434
|
+
# the name-only matches elsewhere are noise, not a shortlist
|
|
435
|
+
return [s for s in scored if s[3] >= 0.7] if scored[0][3] >= 0.7 else scored
|
heimdall/install.py
ADDED
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
# Copyright 2026 The Heimdall Authors
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
"""Register heimdall with a coding agent, globally or in one repo.
|
|
4
|
+
|
|
5
|
+
The config files are written here rather than shelled out to `claude mcp add` /
|
|
6
|
+
`codex mcp add`, so `heimdall install codex` works on a machine that only has Claude,
|
|
7
|
+
and so the whole thing is testable against a fake home directory.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
import os
|
|
14
|
+
import shutil
|
|
15
|
+
import sys
|
|
16
|
+
import tomllib
|
|
17
|
+
from pathlib import Path
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _skill_text() -> str:
|
|
21
|
+
"""Shipped inside the wheel; read from the repo when running off a checkout."""
|
|
22
|
+
here = Path(__file__).resolve()
|
|
23
|
+
for p in (here.with_name("SKILL.md"), here.parents[1] / "skills" / "heimdall" / "SKILL.md"):
|
|
24
|
+
if p.exists():
|
|
25
|
+
return p.read_text(encoding="utf-8")
|
|
26
|
+
raise SystemExit("SKILL.md is missing from this install")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _ephemeral(exe: Path) -> bool:
|
|
30
|
+
"""`uvx` runs out of uv's cache, which uv garbage-collects: a config pointing in
|
|
31
|
+
there rots. `uv tool install` lives under ~/.local/share and does not.
|
|
32
|
+
"""
|
|
33
|
+
# the defaults stay in the list even when UV_CACHE_DIR is set: this shell's setting
|
|
34
|
+
# says nothing about where the binary being inspected was unpacked
|
|
35
|
+
caches = [Path.home() / ".cache" / "uv", Path.home() / "Library" / "Caches" / "uv"]
|
|
36
|
+
if "UV_CACHE_DIR" in os.environ:
|
|
37
|
+
caches.append(Path(os.environ["UV_CACHE_DIR"]))
|
|
38
|
+
return any(exe.is_relative_to(c) for c in caches)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def mcp_command() -> list[str]:
|
|
42
|
+
"""What the agent should spawn to get an MCP server.
|
|
43
|
+
|
|
44
|
+
The copy of heimdall running this command, if it has a durable home -- not whatever
|
|
45
|
+
`heimdall` happens to be first on *this shell's* PATH. Otherwise `uvx`, which costs a
|
|
46
|
+
resolve per start but survives.
|
|
47
|
+
"""
|
|
48
|
+
name = "heimdall.exe" if sys.platform == "win32" else "heimdall"
|
|
49
|
+
exe = Path(sys.executable).with_name(name)
|
|
50
|
+
if not exe.exists():
|
|
51
|
+
exe = Path(shutil.which(name) or "/nonexistent")
|
|
52
|
+
if exe.exists() and not _ephemeral(exe.resolve()):
|
|
53
|
+
return [str(exe), "mcp"]
|
|
54
|
+
return ["uvx", "--from", "heimdall-code", "heimdall", "mcp"]
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _write(path: Path, text: str) -> None:
|
|
58
|
+
path.parent.mkdir(parents=True, exist_ok=True)
|
|
59
|
+
tmp = path.with_suffix(path.suffix + ".tmp")
|
|
60
|
+
tmp.write_text(text, encoding="utf-8")
|
|
61
|
+
os.replace(tmp, path) # never leave a half-written config behind
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def _json_mcp(path: Path, cmd: list[str]) -> str:
|
|
65
|
+
data = {}
|
|
66
|
+
if path.exists():
|
|
67
|
+
try:
|
|
68
|
+
data = json.loads(path.read_text(encoding="utf-8"))
|
|
69
|
+
except ValueError as exc:
|
|
70
|
+
return f"! {path}: not valid JSON ({exc}) -- left alone"
|
|
71
|
+
data.setdefault("mcpServers", {})["heimdall"] = {"command": cmd[0], "args": cmd[1:]}
|
|
72
|
+
_write(path, json.dumps(data, indent=2) + "\n")
|
|
73
|
+
return f"+ MCP server -> {path}"
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _toml_mcp(path: Path, cmd: list[str]) -> str:
|
|
77
|
+
text = path.read_text(encoding="utf-8") if path.exists() else ""
|
|
78
|
+
try:
|
|
79
|
+
existing = tomllib.loads(text)
|
|
80
|
+
except tomllib.TOMLDecodeError as exc:
|
|
81
|
+
return f"! {path}: not valid TOML ({exc}) -- left alone"
|
|
82
|
+
if "heimdall" in existing.get("mcp_servers", {}):
|
|
83
|
+
return f"= MCP server already in {path}"
|
|
84
|
+
# ponytail: append-only, because the stdlib reads TOML and does not write it. An
|
|
85
|
+
# existing entry is reported rather than rewritten -- edit it or delete it by hand.
|
|
86
|
+
args = ", ".join(json.dumps(a) for a in cmd[1:])
|
|
87
|
+
section = f"[mcp_servers.heimdall]\ncommand = {json.dumps(cmd[0])}\nargs = [{args}]\n"
|
|
88
|
+
_write(path, (f"{text.rstrip()}\n\n{section}" if text.strip() else section))
|
|
89
|
+
return f"+ MCP server -> {path}"
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def _skill(path: Path) -> str:
|
|
93
|
+
_write(path, _skill_text())
|
|
94
|
+
return f"+ CLI skill -> {path}"
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
# agent -> (global mcp, project mcp, global skill dir, project skill dir). Claude reads
|
|
98
|
+
# ~/.claude.json and ./.mcp.json; Codex reads ~/.codex/config.toml and, for a trusted
|
|
99
|
+
# project, ./.codex/config.toml. Skills follow each tool's documented directory.
|
|
100
|
+
TARGETS = {
|
|
101
|
+
"claude": (".claude.json", ".mcp.json", ".claude/skills", ".claude/skills"),
|
|
102
|
+
"codex": (".codex/config.toml", ".codex/config.toml", ".agents/skills", ".agents/skills"),
|
|
103
|
+
}
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def install(
|
|
107
|
+
agent: str, *, project: bool = False, cli: bool = False, root: Path | None = None, home: Path | None = None
|
|
108
|
+
) -> list[str]:
|
|
109
|
+
"""Write one agent's config. Returns a line per file touched."""
|
|
110
|
+
root = (root or Path.cwd()).resolve()
|
|
111
|
+
home = home or Path.home()
|
|
112
|
+
base = root if project else home
|
|
113
|
+
out = []
|
|
114
|
+
for a in ["claude", "codex"] if agent == "all" else [agent]:
|
|
115
|
+
gm, pm, gs, ps = TARGETS[a]
|
|
116
|
+
if cli:
|
|
117
|
+
out.append(_skill(base / (ps if project else gs) / "heimdall" / "SKILL.md"))
|
|
118
|
+
else:
|
|
119
|
+
path = base / (pm if project else gm)
|
|
120
|
+
out.append((_toml_mcp if path.suffix == ".toml" else _json_mcp)(path, mcp_command()))
|
|
121
|
+
return out
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
if __name__ == "__main__": # self-check: heimdall install, against a fake home
|
|
125
|
+
import tempfile
|
|
126
|
+
|
|
127
|
+
with tempfile.TemporaryDirectory() as td:
|
|
128
|
+
h, r = Path(td) / "home", Path(td) / "repo"
|
|
129
|
+
(h / ".codex").mkdir(parents=True)
|
|
130
|
+
(h / ".codex" / "config.toml").write_text('model = "gpt-5"\n')
|
|
131
|
+
print(*install("all", home=h, root=r), sep="\n")
|
|
132
|
+
print(*install("all", project=True, home=h, root=r), sep="\n")
|
|
133
|
+
print(*install("claude", cli=True, home=h, root=r), sep="\n")
|
|
134
|
+
|
|
135
|
+
cfg = json.loads((h / ".claude.json").read_text())
|
|
136
|
+
assert cfg["mcpServers"]["heimdall"]["args"][-1] == "mcp", cfg
|
|
137
|
+
codex = tomllib.loads((h / ".codex" / "config.toml").read_text())
|
|
138
|
+
assert codex["model"] == "gpt-5", "clobbered the user's own settings"
|
|
139
|
+
assert codex["mcp_servers"]["heimdall"]["args"][-1] == "mcp", codex
|
|
140
|
+
assert json.loads((r / ".mcp.json").read_text())["mcpServers"]["heimdall"]
|
|
141
|
+
assert tomllib.loads((r / ".codex" / "config.toml").read_text())["mcp_servers"]["heimdall"]
|
|
142
|
+
assert "heimdall index" in (h / ".claude/skills/heimdall/SKILL.md").read_text()
|
|
143
|
+
# a throwaway uvx venv must not be written into a config; a uv tool install must
|
|
144
|
+
assert _ephemeral(Path.home() / ".cache/uv/environments-v2/x/bin/heimdall")
|
|
145
|
+
assert not _ephemeral(Path.home() / ".local/share/uv/tools/heimdall-code/bin/heimdall")
|
|
146
|
+
# rerunning is idempotent, and a corrupt file is reported, not overwritten
|
|
147
|
+
assert install("claude", home=h, root=r) == install("claude", home=h, root=r)
|
|
148
|
+
(h / ".claude.json").write_text("{not json")
|
|
149
|
+
assert install("claude", home=h, root=r)[0].startswith("!")
|
|
150
|
+
assert (h / ".claude.json").read_text() == "{not json"
|
|
151
|
+
print("ok", file=sys.stderr)
|