From b78fe8c597db645ac84d439aca2c7d9c0bc3811a Mon Sep 17 00:00:00 2001 From: garrytan Date: Sat, 17 Oct 2026 17:15:11 +0000 Subject: [PATCH] fix(pregate): keep Python import scanning linear and conservative Review of #3201 on the integrated tree found two quadratic paths on untrusted Python source and three ways a test could be left unselected: - The from-import module spec let two repetitions share a dot, so a failed match on a long run of dots backtracked quadratically. The spec is now dots followed by an optional word-led name. - Dotted literals and imports resolved every prefix with repeated split/join. Names are capped at 32 segments, deeper than any module path. - A literal that also named a repo file skipped its mock.patch reading; both readings now add edges. - A package and a same-named module both map (Python loads the package). - Imports after `;` or a compound statement's `:` are recognised. Co-authored-by: capy-ai[bot] <230910855+capy-ai[bot]@users.noreply.github.com> --- lib/pregate/import-graph.ts | 24 ++++++++++++++---------- test/pregate-selection.test.ts | 24 ++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 10 deletions(-) diff --git a/lib/pregate/import-graph.ts b/lib/pregate/import-graph.ts index 48f022b58..faf4c05c0 100644 --- a/lib/pregate/import-graph.ts +++ b/lib/pregate/import-graph.ts @@ -46,7 +46,10 @@ const JOIN = /\bjoin\(([^()]*)\)/g; const RESOLVE_EXT = ['', '.ts', '.tsx', '.js', '.mjs', '.cjs', '/index.ts', '/index.js']; // One left-to-right pass, so a `#` inside a string and a quote inside a comment both stay inert. const PY_STRING_OR_COMMENT = /("""|''')[\s\S]*?\1|"(?:[^"\\\n]|\\.)*"|'(?:[^'\\\n]|\\.)*'|#[^\n]*/g; -const PY_IMPORT = /^[ \t]*(?:from[ \t]+(\.*[\w.]*)[ \t]+import[ \t]+(\([^)]*\)|[^\n;]+)|import[ \t]+([^\n;]+))/gm; +// A statement starts a line or follows `;` or a compound statement's `:`; the module spec never lets two repetitions share a dot. +const PY_IMPORT = /(?:^|[;:])[ \t\f]*(?:from[ \t]+(\.*(?:\w[\w.]*)?)[ \t]+import[ \t]+(\([^)]*\)|[^\n;]+)|import[ \t]+([^\n;]+))/gm; +// Deeper dotted names are not module paths; the cap keeps prefix resolution linear in the source size. +const PY_MAX_DEPTH = 32; const PY_DOTTED = /^[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)+$/; export function isTestFile(rel: string): boolean { return TEST_FILE_RE.test(rel); } @@ -222,23 +225,24 @@ function pythonImports(rel: string, source: string): string[] { } /** - * A dotted module → its repo file (`a/b.py` or `a/b/__init__.py`) plus the - * package `__init__.py` files Python executes on the way; [] for a module - * outside the repo (stdlib, site-packages). + * A dotted module → its repo files (`a/b/__init__.py`, which Python prefers, + * and `a/b.py`; both when both exist) plus the package `__init__.py` files + * Python executes on the way; [] for a module outside the repo (stdlib, + * site-packages) or deeper than PY_MAX_DEPTH. */ function pythonModuleFiles(files: Set, dotted: string): string[] { const parts = dotted.split('.'); - if (parts.some(p => !p)) return []; + if (parts.length > PY_MAX_DEPTH || parts.some(p => !p)) return []; const p = parts.join('/'); - const file = [`${p}.py`, `${p}/__init__.py`].find(f => files.has(f)); - if (!file) return []; + const found = [`${p}/__init__.py`, `${p}.py`].filter(f => files.has(f)); + if (!found.length) return []; const inits = parts.slice(0, -1).map((_, i) => `${parts.slice(0, i + 1).join('/')}/__init__.py`).filter(f => files.has(f)); - return [file, ...inits]; + return [...found, ...inits]; } /** `'pkg.mod.Class.attr'` (a `mock.patch` target) → the longest dotted prefix that is a repo module. */ function pythonPatchTarget(files: Set, ref: string): string[] { - const parts = ref.split('.'); + const parts = ref.split('.').slice(0, PY_MAX_DEPTH); for (let n = parts.length; n >= 1; n--) { const hit = pythonModuleFiles(files, parts.slice(0, n).join('.')); if (hit.length) return hit; @@ -280,7 +284,7 @@ export function referencesIn(files: Set, dirs: Set, rel: string, for (const m of source.matchAll(QUOTED)) { const targets = resolveRefs(files, dirs, rel, m[1]!, true, config); for (const target of targets) out.push({ target, via: 'literal' }); - if (!targets.length && python && PY_DOTTED.test(m[1]!)) for (const f of pythonPatchTarget(files, m[1]!)) out.push({ target: f, via: 'literal' }); + if (python && PY_DOTTED.test(m[1]!)) for (const f of pythonPatchTarget(files, m[1]!)) out.push({ target: f, via: 'literal' }); } for (const m of source.matchAll(JOIN)) { const segments = [...m[1]!.matchAll(/['"`]([^'"`]+)['"`]/g)].map(s => s[1]!); diff --git a/test/pregate-selection.test.ts b/test/pregate-selection.test.ts index 4b7868060..94ddba92c 100644 --- a/test/pregate-selection.test.ts +++ b/test/pregate-selection.test.ts @@ -113,6 +113,30 @@ describe('import graph (static imports ∪ repo-path literals ∪ bin convention 'pkg/__init__.py', 'pkg/core.py', 'pkg/multi.py', 'pkg/patched.py', 'pkg/sub/__init__.py', 'pkg/sub/leaf.py', 'pkg/sub/sibling.py', 'pkg/util.py', ]); }); + + test('python: a package and a same-named module both map, statements after ; and : map, a path literal does not hide its patch target, and adversarial dotted input stays linear', () => { + const repo = fixture(); + commit(repo, { + 'app.py': 'A = 1\n', + 'app/__init__.py': '', + 'pkg/__init__.py': '', + 'pkg/core.py': 'C = 1\n', + 'pkg/patched.py': 'class Thing:\n pass\n', + 'pkg.patched.Thing.run': 'a file whose name is also a patch target\n', + 'tests/test_app.py': 'import app\n', + 'tests/test_semicolon.py': 'import os; import pkg.core\n', + 'tests/test_inline.py': 'def test_core():\n if True: from pkg.core import C\n', + 'tests/test_patch.py': "from unittest import mock\n\n@mock.patch('pkg.patched.Thing.run')\ndef test_run(run):\n pass\n", + 'tests/test_adversarial.py': `from ${'.'.repeat(100_000)}x\nX = '${Array.from({ length: 20_000 }, () => 'pkg').join('.')}'\n`, + }); + const g = buildGraph(repo.work); + expect(testsFor(g, 'app/__init__.py').tests).toEqual(['tests/test_app.py']); + expect(testsFor(g, 'app.py').tests).toEqual(['tests/test_app.py']); + expect(testsFor(g, 'pkg/core.py').tests).toEqual(['tests/test_inline.py', 'tests/test_semicolon.py']); + expect(testsFor(g, 'pkg/patched.py').tests).toEqual(['tests/test_patch.py']); + expect(testsFor(g, 'pkg.patched.Thing.run').tests).toEqual(['tests/test_patch.py']); + expect(testsFor(g, 'pkg/__init__.py').tests).toEqual(['tests/test_adversarial.py', 'tests/test_inline.py', 'tests/test_patch.py', 'tests/test_semicolon.py']); + }); }); describe('literals (pure)', () => {