argot-engine 0.2.0__tar.gz → 0.2.2__tar.gz

{argot_engine-0.2.0 → argot_engine-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: argot-engine
-Version: 0.2.0
+Version: 0.2.2
 Requires-Python: >=3.11
 Requires-Dist: pygit2==1.19.2
 Requires-Dist: scikit-learn>=1.5.0

{argot_engine-0.2.0 → argot_engine-0.2.2}/argot/check.py

@@ -2,13 +2,15 @@ from __future__ import annotations
 
 import argparse
 import sys
+from collections.abc import Iterator
 from pathlib import Path
+from typing import Any
 
 import joblib  # type: ignore[import-untyped]
 import pygit2
 import torch
 
-from argot.git_walk import walk_commits
+from argot.git_walk import SUPPORTED_EXTENSIONS, _extension, walk_commits
 from argot.jepa.encoder import TokenEncoder
 from argot.jepa.model import JEPAArgot
 from argot.jepa.predictor import ArgotPredictor
@@ -37,42 +39,43 @@ def _resolve_shas(repo: pygit2.Repository, ref: str) -> set[str]:
     return shas
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Check code surprise with argot JEPA model")
-    parser.add_argument("repo_path")
-    parser.add_argument("ref")
-    parser.add_argument("--model", default=".argot/model.pkl")
-    parser.add_argument("--threshold", type=float, default=0.5)
-    args = parser.parse_args()
-
-    model_path = Path(args.model)
-    if not model_path.exists():
-        print(f"error: model not found at {model_path}", file=sys.stderr)
-        sys.exit(2)
-
-    bundle = joblib.load(model_path)
-    vectorizer = bundle["vectorizer"]
-    embed_dim: int = bundle["embed_dim"]
-    input_dim: int = bundle["input_dim"]
-
-    encoder = TokenEncoder(input_dim, embed_dim)
-    encoder.load_state_dict(bundle["encoder_state"])
-    predictor = ArgotPredictor(embed_dim=embed_dim)
-    predictor.load_state_dict(bundle["predictor_state"])
-    model = JEPAArgot(encoder, predictor)
-    model.eval()
-
-    repo = pygit2.Repository(args.repo_path)
-    shas = _resolve_shas(repo, args.ref)
-    if not shas:
-        print("No commits found in range", file=sys.stderr)
-        sys.exit(0)
-
+def _workdir_patches(
+    repo_path: str,
+) -> Iterator[tuple[str, bytes, list[pygit2.DiffHunk]]]:
+    """Yield (file_path, content, hunks) for uncommitted changes vs HEAD."""
+    repo = pygit2.Repository(repo_path)
+    head_oid = repo.revparse_single("HEAD").id
+    diff = repo.diff(a=head_oid)
+    diff.find_similar()
+    workdir = Path(repo.workdir)
+    for patch in diff:
+        if patch is None:
+            continue
+        file_path = patch.delta.new_file.path
+        if _extension(file_path) not in SUPPORTED_EXTENSIONS:
+            continue
+        hunks = list(patch.hunks)
+        if not hunks:
+            continue
+        full_path = workdir / file_path
+        if not full_path.exists():
+            continue
+        yield file_path, full_path.read_bytes(), hunks
+
+
+def _score_patches(
+    patches: Iterator[tuple[str, bytes, list[pygit2.DiffHunk]]],
+    vectorizer: Any,
+    model: JEPAArgot,
+    label: str,
+) -> tuple[list[tuple[float, str, int, str]], int]:
+    """Score hunk patches; returns (results, total_hunk_count)."""
+    context_lines = 50
     results: list[tuple[float, str, int, str]] = []
+    hunk_count = 0
 
-    context_lines = 50
     with torch.no_grad():
-        for commit, file_path, post_blob, hunks in walk_commits(args.repo_path, shas):
+        for file_path, post_blob, hunks in patches:
             lang = language_for_path(file_path)
             if lang is None:
                 continue
@@ -82,6 +85,7 @@ def main() -> None:
                 continue
 
             for hunk in hunks:
+                hunk_count += 1
                 hunk_start = hunk.new_start - 1
                 hunk_end = hunk_start + hunk.new_lines
                 if hunk_start < 0 or hunk_end > len(source_lines):
@@ -102,17 +106,89 @@ def main() -> None:
                 )
 
                 score = model.surprise(ctx_vec, hunk_vec).item()
-                results.append((score, file_path, hunk.new_start, str(commit.id)[:8]))
+                results.append((score, file_path, hunk.new_start, label))
+
+    return results, hunk_count
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Check code surprise with argot JEPA model")
+    parser.add_argument("repo_path")
+    parser.add_argument("ref", nargs="?", default="")
+    parser.add_argument("--model", default=".argot/model.pkl")
+    parser.add_argument("--threshold", type=float, default=0.5)
+    args = parser.parse_args()
+
+    model_path = Path(args.model)
+    if not model_path.exists():
+        print(f"error: model not found at {model_path}", file=sys.stderr)
+        sys.exit(2)
+
+    bundle = joblib.load(model_path)
+    vectorizer = bundle["vectorizer"]
+    embed_dim: int = bundle["embed_dim"]
+    input_dim: int = bundle["input_dim"]
+
+    encoder = TokenEncoder(input_dim, embed_dim)
+    encoder.load_state_dict(bundle["encoder_state"])
+    predictor = ArgotPredictor(embed_dim=embed_dim)
+    predictor.load_state_dict(bundle["predictor_state"])
+    model = JEPAArgot(encoder, predictor)
+    model.eval()
+
+    if args.ref == "":
+        patches: Iterator[tuple[str, bytes, list[pygit2.DiffHunk]]] = _workdir_patches(
+            args.repo_path
+        )
+        context_label = "workdir"
+        commit_info = "working tree"
+    else:
+        repo = pygit2.Repository(args.repo_path)
+        shas = _resolve_shas(repo, args.ref)
+        if not shas:
+            print("No commits found in range", file=sys.stderr)
+            sys.exit(0)
+
+        def _committed_patches() -> Iterator[tuple[str, bytes, list[pygit2.DiffHunk]]]:
+            for _commit, file_path, post_blob, hunks in walk_commits(args.repo_path, shas):
+                yield file_path, post_blob, hunks
+
+        patches = _committed_patches()
+        context_label = args.ref
+        commit_info = f"{len(shas)} commit(s)"
+
+    results, hunk_count = _score_patches(patches, vectorizer, model, context_label)
 
     if not results:
-        print("No hunks found in range")
+        if hunk_count == 0:
+            exts = " ".join(sorted(SUPPORTED_EXTENSIONS))
+            print(
+                f"No changes to supported files found ({commit_info} scanned).\n"
+                f"Supported extensions: {exts}"
+            )
+            if args.ref != "":
+                print("Try a wider range, e.g.: argot check HEAD~20..HEAD")
+        else:
+            print(
+                f"All {hunk_count} hunk(s) scored below threshold {args.threshold:.2f}"
+                " — looks clean."
+            )
         sys.exit(0)
 
     results.sort(key=lambda r: r[0], reverse=True)
 
-    print(f"{'SURPRISE':>9}  {'FILE':<48}  {'LINE':>5}  COMMIT")
-    for score, fp, line, sha in results:
-        print(f"{score:>9.4f}  {fp:<48}  {line:>5}  {sha}")
+    t = args.threshold
+    print(f"{'SURPRISE':>9}  {'TAG':<10}  {'FILE':<48}  {'LINE':>5}  REF")
+    for score, fp, line, ref in results:
+        if score <= t:
+            tag = "ok"
+        elif score <= t + 0.3:
+            tag = "unusual"
+        elif score <= t + 0.6:
+            tag = "suspicious"
+        else:
+            tag = "foreign"
+        print(f"{score:>9.4f}  {tag:<10}  {fp:<48}  {line:>5}  {ref}")
 
     if any(s > args.threshold for s, *_ in results):
         sys.exit(1)

argot_engine-0.2.2/argot/tests/test_check.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import cast
+
+import pygit2
+
+from argot.check import _resolve_shas, _workdir_patches
+
+
+def _make_repo(tmp_path: Path, files: dict[str, str]) -> pygit2.Repository:
+    """Create a repo with a single commit containing the given files."""
+    repo = pygit2.init_repository(str(tmp_path))
+    sig = pygit2.Signature("Test", "test@example.com")
+    for name, content in files.items():
+        (tmp_path / name).write_text(content)
+        repo.index.add(name)
+    repo.index.write()
+    tree = repo.index.write_tree()
+    repo.create_commit("refs/heads/main", sig, sig, "init", tree, [])
+    repo.set_head("refs/heads/main")
+    return repo
+
+
+def _make_two_commit_repo(tmp_path: Path) -> pygit2.Repository:
+    repo = pygit2.init_repository(str(tmp_path))
+    sig = pygit2.Signature("Test", "test@example.com")
+    f = tmp_path / "main.py"
+    f.write_text("x = 1\n")
+    repo.index.add("main.py")
+    repo.index.write()
+    tree1 = repo.index.write_tree()
+    c1 = repo.create_commit("refs/heads/main", sig, sig, "first", tree1, [])
+    f.write_text("x = 2\n")
+    repo.index.add("main.py")
+    repo.index.write()
+    tree2 = repo.index.write_tree()
+    repo.create_commit("refs/heads/main", sig, sig, "second", tree2, [c1])
+    return repo
+
+
+def test_resolve_shas_range(tmp_path: Path) -> None:
+    repo = _make_two_commit_repo(tmp_path)
+    head_oid = repo.references["refs/heads/main"].target
+    commit = cast(pygit2.Commit, repo.get(head_oid))
+    parent_oid = commit.parents[0].id
+
+    shas = _resolve_shas(repo, f"{parent_oid}..refs/heads/main")
+    assert str(head_oid) in shas
+    assert str(parent_oid) not in shas
+
+
+def test_resolve_shas_bare_ref(tmp_path: Path) -> None:
+    repo = _make_two_commit_repo(tmp_path)
+    head_oid = str(repo.references["refs/heads/main"].target)
+    shas = _resolve_shas(repo, "refs/heads/main")
+    assert head_oid in shas
+
+
+def test_workdir_patches_detects_modification(tmp_path: Path) -> None:
+    _make_repo(tmp_path, {"main.py": "x = 1\n"})
+    (tmp_path / "main.py").write_text("x = 1\ny = 2\n")
+
+    patches = list(_workdir_patches(str(tmp_path)))
+    assert len(patches) == 1
+    file_path, content, hunks = patches[0]
+    assert file_path == "main.py"
+    assert b"y = 2" in content
+    assert len(hunks) > 0
+
+
+def test_workdir_patches_ignores_unsupported_extension(tmp_path: Path) -> None:
+    _make_repo(tmp_path, {"config.json": "{}\n"})
+    (tmp_path / "config.json").write_text('{"key": "value"}\n')
+
+    patches = list(_workdir_patches(str(tmp_path)))
+    assert patches == []
+
+
+def test_workdir_patches_ignores_deleted_files(tmp_path: Path) -> None:
+    _make_repo(tmp_path, {"main.py": "x = 1\n"})
+    (tmp_path / "main.py").unlink()
+
+    patches = list(_workdir_patches(str(tmp_path)))
+    assert patches == []
+
+
+def test_workdir_patches_no_changes(tmp_path: Path) -> None:
+    _make_repo(tmp_path, {"main.py": "x = 1\n"})
+
+    patches = list(_workdir_patches(str(tmp_path)))
+    assert patches == []

{argot_engine-0.2.0 → argot_engine-0.2.2}/argot_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: argot-engine
-Version: 0.2.0
+Version: 0.2.2
 Requires-Python: >=3.11
 Requires-Dist: pygit2==1.19.2
 Requires-Dist: scikit-learn>=1.5.0

{argot_engine-0.2.0 → argot_engine-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "argot-engine"
-version = "0.2.0"
+version = "0.2.2"
 requires-python = ">=3.11"
 dependencies = [
     "pygit2==1.19.2",

argot_engine-0.2.0/argot/tests/test_check.py

@@ -1,43 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-from typing import cast
-
-import pygit2
-
-from argot.check import _resolve_shas
-
-
-def _make_two_commit_repo(tmp_path: Path) -> pygit2.Repository:
-    repo = pygit2.init_repository(str(tmp_path))
-    sig = pygit2.Signature("Test", "test@example.com")
-    f = tmp_path / "main.py"
-    f.write_text("x = 1\n")
-    repo.index.add("main.py")
-    repo.index.write()
-    tree1 = repo.index.write_tree()
-    c1 = repo.create_commit("refs/heads/main", sig, sig, "first", tree1, [])
-    f.write_text("x = 2\n")
-    repo.index.add("main.py")
-    repo.index.write()
-    tree2 = repo.index.write_tree()
-    repo.create_commit("refs/heads/main", sig, sig, "second", tree2, [c1])
-    return repo
-
-
-def test_resolve_shas_range(tmp_path: Path) -> None:
-    repo = _make_two_commit_repo(tmp_path)
-    head_oid = repo.references["refs/heads/main"].target
-    commit = cast(pygit2.Commit, repo.get(head_oid))
-    parent_oid = commit.parents[0].id
-
-    shas = _resolve_shas(repo, f"{parent_oid}..refs/heads/main")
-    assert str(head_oid) in shas
-    assert str(parent_oid) not in shas
-
-
-def test_resolve_shas_bare_ref(tmp_path: Path) -> None:
-    repo = _make_two_commit_repo(tmp_path)
-    head_oid = str(repo.references["refs/heads/main"].target)
-    shas = _resolve_shas(repo, "refs/heads/main")
-    assert head_oid in shas