Parse extensionless Python and Node scripts by their shebang

README.md

@@ -89,7 +89,8 @@ list.
 
 The summary and compact file list come first, followed by selected source with
 line references, then detailed declaration and call locations. Python, TypeScript/JavaScript, Go and Rust support declaration
-parsing; other text uses a fallback. The output is evidence for the agent to use,
+parsing, with extensionless Python and Node scripts recognised by their shebang;
+other text uses a fallback. The output is evidence for the agent to use,
 not a generated answer or a guarantee that every relevant file was found.
 [See a recorded output example](specs/done/jevgrep/assets/stdout-example.txt).
 

packages/core/src/call-context.ts

@@ -1,4 +1,5 @@
 import { runParser } from "./parser";
+import { resolveLanguage } from "./source";
 import type { Snapshot } from "./filesystem";
 import type { FileEvidence, Range } from "./types";
 type Call = Range & {
@@ -14,7 +15,7 @@ export async function localCallContext(
 ) {
   const original = file.presentationExcerpts ?? file.excerpts;
   if (
-    !/\.pyi?$/.test(snapshot.path) ||
+    resolveLanguage(snapshot.path, snapshot.source) !== "python" ||
     Buffer.byteLength(snapshot.source) > 1_000_000 ||
     file.sourceOmitted ||
     original.some((e) => e.partial || e.sourceByteStart !== undefined)

packages/core/src/parser-helpers.mjs

@@ -432,7 +432,7 @@ export async function execute(helper, input) {
       data.query.normalize("NFKC").match(/[_\p{ID_Start}][_\p{ID_Continue}]*/gu) ?? [],
     );
     let matches = [];
-    if (/\.pyi?$/.test(data.path) && Buffer.byteLength(data.text) > data.budget)
+    if (data.language === "python" && Buffer.byteLength(data.text) > data.budget)
       matches = await parse(data.text, (root) => previewMatches(root, tokens));
     return preview(data, matches);
   }

packages/core/src/retrieve.ts

@@ -1,7 +1,7 @@
 import { basename, extname } from "node:path";
 import { createFilesystem, type Snapshot, type DirectoryEntry } from "./filesystem";
 import { type EvaluationRequest, EvaluationFailure } from "./evaluator";
-import { inspect, pythonPreview, sourceForUnit, splitSource } from "./source";
+import { inspect, pythonPreview, resolveLanguage, sourceForUnit, splitSource } from "./source";
 import {
   navigationRequest,
   fileAssessmentRequest,
@@ -304,6 +304,7 @@ export async function retrieve(input: SearchInput, evaluator: Evaluator): Promis
     return buffered({ ...item, childPreview: preview }, sources);
   }
   async function previewFile(source: Snapshot): Promise<FilePreview> {
+    const language = resolveLanguage(source.path, source.source);
     const bytes = Buffer.from(source.source);
     let text = new TextDecoder("utf8", { fatal: true }).decode(bytes.subarray(0, 16384), {
       stream: bytes.length > 16384,
@@ -326,7 +327,7 @@ export async function retrieve(input: SearchInput, evaluator: Evaluator): Promis
       declarations: [],
       declarationIndexTruncated: false,
     };
-    if (truncated && /\.pyi?$/.test(source.path) && bytes.length <= 1_000_000) {
+    if (truncated && language === "python" && bytes.length <= 1_000_000) {
       const sampled = await pythonPreview(source, input.query, 16384, input.signal);
       if (
         sampled?.truncated &&
@@ -339,11 +340,7 @@ export async function retrieve(input: SearchInput, evaluator: Evaluator): Promis
         preview.range = "sampled source ranges";
       }
     }
-    if (
-      truncated &&
-      /\.(?:pyi?|[cm]?[jt]s|[jt]sx|go|rs)$/.test(source.path) &&
-      bytes.length <= 1_000_000
-    ) {
+    if (truncated && language !== undefined && bytes.length <= 1_000_000) {
       const syntax = await inspect(source, {
         signal: input.signal,
         maxUnitBytes: Math.max(4, bytes.length),

packages/core/src/source.ts

@@ -1,4 +1,5 @@
 import ts from "typescript";
+import { basename, extname } from "node:path";
 import { runParser } from "./parser.ts";
 
 export type Range = { startLine: number; endLine: number };
@@ -19,6 +20,29 @@ export type Inspection = {
   fallback?: "unsupported" | "syntax" | "size";
 };
 
+export function resolveLanguage(
+  path: string,
+  source: string,
+): "python" | "go" | "rust" | "TS" | "TSX" | "JS" | "JSX" | undefined {
+  if (/\.pyi?$/.test(path)) return "python";
+  if (path.endsWith(".tsx")) return "TSX";
+  if (path.endsWith(".jsx")) return "JSX";
+  if (/\.[cm]?js$/.test(path)) return "JS";
+  if (/\.[cm]?ts$/.test(path)) return "TS";
+  if (path.endsWith(".go")) return "go";
+  if (path.endsWith(".rs")) return "rust";
+  if (extname(path) !== "" || !source.startsWith("#!")) return;
+  const [program = "", ...args] = source.split("\n", 1)[0]!.slice(2).trim().split(/\s+/);
+  // env options, the operand of -u/-C/-P/-a and NAME=VALUE assignments come before the command.
+  const command = args.find(
+    (arg, index) =>
+      !arg.startsWith("-") && !arg.includes("=") && !/^-[uCPa]$/.test(args[index - 1] ?? ""),
+  );
+  const interpreter = basename(basename(program) === "env" ? (command ?? "") : program);
+  if (/^python(\d+(\.\d+)?)?$/.test(interpreter)) return "python";
+  if (interpreter === "node" || interpreter === "nodejs") return "JS";
+}
+
 type SourceText = { bytes: Buffer; offsets: number[]; lineCount: number };
 function sourceText(source: string): SourceText {
   const lines = source.split("\n"),
@@ -82,15 +106,17 @@ export async function inspect(
       "Source bounds must be integers; unit bytes at least 4 and parse bytes positive",
     );
   const { source, path } = snapshot;
+  const language = resolveLanguage(path, source);
   const text = sourceText(source);
   // Context windows use conservative whole-line Python comments, including inside multiline strings.
-  const pythonComments = /\.pyi?$/.test(path)
-    ? source
-        .split("\n")
-        .flatMap((line, index) =>
-          /^\s*#/.test(line) ? [{ startLine: index + 1, endLine: index + 1 }] : [],
-        )
-    : [];
+  const pythonComments =
+    language === "python"
+      ? source
+          .split("\n")
+          .flatMap((line, index) =>
+            /^\s*#/.test(line) ? [{ startLine: index + 1, endLine: index + 1 }] : [],
+          )
+      : [];
   const fallback = (reason: Inspection["fallback"]): Inspection => ({
     mode: "text",
     fallback: reason,
@@ -104,7 +130,7 @@ export async function inspect(
     comments: Range[] = [],
     mode: Inspection["mode"];
   let syntaxFallback = false;
-  if (/\.pyi?$/.test(path)) {
+  if (language === "python") {
     // A missing/incompatible packaged parser is a setup failure, never syntax fallback.
     const parsed = await runParser<Array<Range & { name: string; ownerHeaders: Range[] }>>(
       "inspect",
@@ -119,15 +145,14 @@ export async function inspect(
     }));
     comments = pythonComments;
     mode = "python";
-  } else if (/\.(?:[cm]?[jt]s|[jt]sx)$/.test(path)) {
-    const kind = path.endsWith(".tsx")
-      ? ts.ScriptKind.TSX
-      : path.endsWith(".jsx")
-        ? ts.ScriptKind.JSX
-        : /\.[cm]?js$/.test(path)
-          ? ts.ScriptKind.JS
-          : ts.ScriptKind.TS;
-    const file = ts.createSourceFile(path, source, ts.ScriptTarget.Latest, true, kind);
+  } else if (language === "TS" || language === "TSX" || language === "JS" || language === "JSX") {
+    const file = ts.createSourceFile(
+      path,
+      source,
+      ts.ScriptTarget.Latest,
+      true,
+      ts.ScriptKind[language],
+    );
     syntaxFallback = !!(file as ts.SourceFile & { parseDiagnostics?: unknown[] }).parseDiagnostics
       ?.length;
     const line = (position: number) => file.getLineAndCharacterOfPosition(position).line + 1;
@@ -167,16 +192,16 @@ export async function inspect(
     };
     visit(file);
     mode = "typescript";
-  } else if (/\.(go|rs)$/.test(path)) {
+  } else if (language === "go" || language === "rust") {
     const parsed = await runParser<{ units: typeof units; comments: Range[] }>(
       "declarations",
-      JSON.stringify({ source, language: path.endsWith(".go") ? "go" : "rust" }),
+      JSON.stringify({ source, language }),
       options.signal,
     );
     if (!parsed) return fallback("syntax");
     units = parsed.units;
     comments = parsed.comments;
-    mode = path.endsWith(".go") ? "go" : "rust";
+    mode = language;
   } else return fallback("unsupported");
   comments = [...new Map(comments.map((r) => [`${r.startLine}:${r.endLine}`, r])).values()].sort(
     (a, b) => a.startLine - b.startLine,
@@ -233,7 +258,11 @@ export async function pythonNeighborhood(
   selected: Range[],
   signal?: AbortSignal,
 ): Promise<Range[]> {
-  if (!/\.pyi?$/.test(snapshot.path) || Buffer.byteLength(snapshot.source) > 1_000_000) return [];
+  if (
+    resolveLanguage(snapshot.path, snapshot.source) !== "python" ||
+    Buffer.byteLength(snapshot.source) > 1_000_000
+  )
+    return [];
   return (
     (await runParser<Range[]>(
       "neighborhood",
@@ -275,7 +304,12 @@ export async function pythonPreview(
     throw new Error("Preview allowance must be at least 256 bytes");
   const result = await runParser<SourcePreview>(
     "preview",
-    JSON.stringify({ text: snapshot.source, path: snapshot.path, query, budget }),
+    JSON.stringify({
+      text: snapshot.source,
+      language: resolveLanguage(snapshot.path, snapshot.source),
+      query,
+      budget,
+    }),
     signal,
   );
   if (result?.truncated) {

packages/core/src/test-body-selection.ts

@@ -1,4 +1,4 @@
-import { inspect } from "./source";
+import { inspect, resolveLanguage } from "./source";
 import type { Snapshot } from "./filesystem";
 import type { Evaluator, FileEvidence, Range } from "./types";
 import { EvaluationFailure } from "./evaluator";
@@ -26,7 +26,7 @@ export async function selectTestBodies(
     if (
       !file.roles.includes("test") ||
       file.sourceOmitted ||
-      !/\.pyi?$/.test(snapshot.path) ||
+      resolveLanguage(snapshot.path, snapshot.source) !== "python" ||
       Buffer.byteLength(snapshot.source) > 1_000_000 ||
       shown.some((e) => e.partial || e.sourceByteStart !== undefined)
     )

test/parser/parity.test.ts

@@ -34,7 +34,13 @@ test("Python declaration, preview, neighborhood and inherited-call outputs match
             ? source
             : JSON.stringify(
                 helper === "preview"
-                  ? { path: "x.py", text: source, query: "target first K", budget: 600 }
+                  ? {
+                      path: "x.py",
+                      language: "python",
+                      text: source,
+                      query: "target first K",
+                      budget: 600,
+                    }
                   : {
                       source,
                       ranges: [

test/parser/source.test.ts

@@ -1,6 +1,44 @@
 import { test } from "node:test";
 import assert from "node:assert/strict";
-import { inspect } from "../../packages/core/src/source.ts";
+import { inspect, pythonNeighborhood } from "../../packages/core/src/source.ts";
+
+test("extensionless Python shebangs preserve declarations, comments and neighboring source", async () => {
+  for (const shebang of [
+    "#!/usr/bin/env python3",
+    "#!/usr/bin/env -S python3 -u",
+    "#!/usr/bin/env -S PYTHONUNBUFFERED=1 python3",
+    "#!/usr/bin/env -S -u PYTHONPATH python3",
+    "#! /usr/bin/python3.12\r",
+  ]) {
+    const source =
+      shebang +
+      '\n# module\nclass Example:\n    """context"""\n    def first(self):\n        return "é"\n    @decorator\n    def middle(self):\n        return self.first()\n    def last(self):\n        pass\n';
+    const snapshot = { path: "bin/tool", contentHash: "fixture", source };
+    const withExtension = { ...snapshot, path: "bin/tool.py" };
+    const expected = await inspect(withExtension);
+    assert.equal(expected.mode, "python");
+    assert.deepEqual(await inspect(snapshot), expected, shebang);
+    const selected = [{ startLine: 8, endLine: 9 }];
+    const neighbors = await pythonNeighborhood(withExtension, selected);
+    assert.ok(neighbors.some((range) => range.startLine === 5 && range.endLine === 6));
+    assert.deepEqual(await pythonNeighborhood(snapshot, selected), neighbors, shebang);
+  }
+});
+
+test("extensionless Node shebangs preserve JavaScript declarations and selected source", async () => {
+  for (const shebang of [
+    "#!/usr/bin/env node",
+    "#!/usr/bin/nodejs",
+    "#!/usr/bin/env -S NODE_ENV=production node",
+  ]) {
+    const source = shebang + '\n// module\nfunction target() {\n  return "🙂";\n}\n';
+    const snapshot = { path: "bin/tool", contentHash: "fixture", source };
+    const expected = await inspect({ ...snapshot, path: "bin/tool.js" });
+    assert.equal(expected.mode, "typescript");
+    assert.deepEqual(await inspect(snapshot), expected, shebang);
+    assert.equal(expected.units[0]!.name, "target");
+  }
+});
 
 test("Python preserves decorators, class context and nested declarations", async () => {
   const source =
@@ -58,6 +96,11 @@ test("syntax errors and unsupported source fall back without losing source lines
     ["bad.py", "def broken(:\n  return 2\n", "syntax"],
     ["bad.ts", "const = ;", "syntax"],
     ["readme.md", "éé\r\nnext\r\nlast", "unsupported"],
+    ["tool", "def target():\n    pass\n", "unsupported"],
+    ["tool", "#!/bin/bash\ndef target():\n    pass\n", "unsupported"],
+    ["tool", "\n#!/usr/bin/env python3\ndef target():\n    pass\n", "unsupported"],
+    ["tool.sh", "#!/usr/bin/env python3\ndef target():\n    pass\n", "unsupported"],
+    ["tool.js", "#!/usr/bin/env python3\ndef target():\n    pass\n", "syntax"],
   ]) {
     const result = await inspect(
       { path: path!, source: source!, contentHash: "fixture" },

test/retrieval.test.ts

@@ -1,13 +1,89 @@
 import { renderResult } from "../apps/cli/src/render";
 import { routeProviderFetch } from "./fixtures/provider-route.mjs";
-import { expect } from "bun:test";
+import { expect, test } from "bun:test";
 import { mkdtemp, mkdir, writeFile, rm } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { testIfDocker } from "./helpers/docker";
 import { retrieve } from "../packages/core/src/retrieve";
 import { createEvaluator, EvaluationFailure } from "../packages/core/src/evaluator";
 
+test("extensionless scripts preserve extension-based provider contents and returned evidence", async () => {
+  for (const [extension, shebang, body] of [
+    [
+      ".py",
+      "#!/usr/bin/env -S python3 -u",
+      'class Base:\n    def helper(self):\n        return "🙂"\nclass Child(Base):\n    def test_target(self):\n        assert self.helper() == "🙂"\n',
+    ],
+    [".js", "#!/usr/bin/env node", 'function target() {\n  return "🙂";\n}\n'],
+  ]) {
+    const root = await mkdtemp(join(tmpdir(), "jg-shebang-"));
+    const padding = `${extension === ".py" ? "#" : "//"} ${"padding ".repeat(12)}\n`;
+    const source = `${shebang}\n${padding.repeat(200)}${body}`;
+    const runs = [];
+    try {
+      for (const path of [`tool${extension}`, "tool"]) {
+        const requests: string[] = [];
+        await writeFile(join(root, path), source);
+        const signal = new AbortController().signal;
+        const result = await retrieve(
+          { root, query: "test_target target", signal },
+          createEvaluator({
+            provider: "vercel",
+            apiKey: "fixture",
+            signal,
+            fetch: async (_input, init) => {
+              const request = String(init?.body);
+              requests.push(request);
+              const body = JSON.parse(request) as {
+                questions: Record<string, { instructions: string }>;
+              };
+              return Response.json({
+                answers: Object.fromEntries(
+                  Object.entries(body.questions).map(([id, question]) => [
+                    id,
+                    {
+                      type: "noul",
+                      noul: question.instructions.includes("(Base.helper,") ? 0.1 : 0.95,
+                    },
+                  ]),
+                ),
+              });
+            },
+          }),
+        );
+        expect(result.status).toBe("complete");
+        expect(requests.join("\n")).toContain(`"extension":"${path === "tool" ? "" : extension}"`);
+        runs.push({
+          requests: requests.map((request) =>
+            request
+              .replaceAll(`tool${extension}`, "tool")
+              .replaceAll(`"extension":"${extension}"`, '"extension":""'),
+          ),
+          files: JSON.parse(JSON.stringify(result.files).replaceAll(`tool${extension}`, "tool")),
+        });
+        await rm(join(root, path));
+      }
+      expect(runs[0]!.requests.join("\n")).toContain('"name":');
+      if (extension === ".py") {
+        expect(runs[0]!.requests.join("\n")).toContain('"range":"sampled source ranges"');
+        expect(runs[0]!.requests.join("\n")).toContain('"candidates":');
+        expect(runs[0]!.files[0].callLeads).toEqual([
+          {
+            caller: "Child.test_target",
+            name: "Base.helper",
+            range: { startLine: 203, endLine: 204 },
+            unknownEarlierBases: [],
+          },
+        ]);
+      }
+      expect(runs[1]).toEqual(runs[0]);
+    } finally {
+      await rm(root, { recursive: true, force: true });
+    }
+  }
+});
+
 testIfDocker(
   "hierarchical retrieval returns source from multiple files without uploading excluded data",
   async () => {