From ef319975de3d93034b83f7edeb5cb45cdde316f9 Mon Sep 17 00:00:00 2001 From: Chirag Nighut Date: Wed, 18 Mar 2026 18:35:10 +0530 Subject: [PATCH] feat: add Python enumerate() for-loop support with nested tuple patterns (#356) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Handle `for i, k, v in enumerate(d.items())` — flat pattern - Handle `for i, (k, v) in enumerate(d.items())` — nested tuple_pattern - Handle `for (k, v) in enumerate(users)` — parenthesized tuple as top-level Extract helper functions for cleaner code: - `extractMethodCall()` — deduplicate method call parsing - `collectPatternIdentifiers()` — recursively collect identifiers from patterns Add unit tests for TypeEnv and integration tests verifying CALLS edges. Made-with: Cursor Co-authored-by: chirag-nighut --- .../core/ingestion/type-extractors/python.ts | 73 +++++++++++++++---- .../python-enumerate-loop/app.py | 20 +++++ .../python-enumerate-loop/repo.py | 3 + .../python-enumerate-loop/user.py | 6 ++ .../test/integration/resolvers/python.test.ts | 57 +++++++++++++++ gitnexus/test/unit/type-env.test.ts | 44 +++++++++++ 6 files changed, 188 insertions(+), 15 deletions(-) create mode 100644 gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/app.py create mode 100644 gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/repo.py create mode 100644 gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/user.py diff --git a/gitnexus/src/core/ingestion/type-extractors/python.ts b/gitnexus/src/core/ingestion/type-extractors/python.ts index 14fc1602d3..a202bdaf29 100644 --- a/gitnexus/src/core/ingestion/type-extractors/python.ts +++ b/gitnexus/src/core/ingestion/type-extractors/python.ts @@ -219,6 +219,37 @@ const findPyParamElementType = (iterableName: string, startNode: SyntaxNode, pos return undefined; }; +/** + * Extracts iterableName and methodName from a call expression like `data.items()`. + * Returns undefined if the call doesn't match the expected pattern. + */ +const extractMethodCall = (callNode: SyntaxNode): { iterableName: string; methodName?: string } | undefined => { + const fn = callNode.childForFieldName('function'); + if (fn?.type !== 'attribute') return undefined; + const obj = fn.firstNamedChild; + if (obj?.type !== 'identifier') return undefined; + const method = fn.lastNamedChild; + const methodName = (method?.type === 'identifier' && method !== obj) ? method.text : undefined; + return { iterableName: obj.text, methodName }; +}; + +/** + * Collects all identifier nodes from a pattern, descending into nested tuple_patterns. + * For `i, (k, v)` returns [i, k, v]. For `key, value` returns [key, value]. + */ +const collectPatternIdentifiers = (pattern: SyntaxNode): SyntaxNode[] => { + const vars: SyntaxNode[] = []; + for (let i = 0; i < pattern.namedChildCount; i++) { + const child = pattern.namedChild(i); + if (child?.type === 'identifier') { + vars.push(child); + } else if (child?.type === 'tuple_pattern') { + vars.push(...collectPatternIdentifiers(child)); + } + } + return vars; +}; + /** * Python: for user in users: where users has a known container type annotation. * @@ -228,32 +259,43 @@ const findPyParamElementType = (iterableName: string, startNode: SyntaxNode, pos * 1. declarationTypeNodes — raw type annotation AST node (covers stored container types) * 2. scopeEnv string — extractElementTypeFromString on the stored type * 3. AST walk — walks up to the enclosing function's parameters to read List[User] directly + * + * Also handles `enumerate(iterable)` — unwraps the outer call and skips the integer + * index variable so the value variable still resolves to the element type. */ const extractForLoopBinding: ForLoopExtractor = (node, { scopeEnv, declarationTypeNodes, scope, returnTypeLookup }): void => { if (node.type !== 'for_statement') return; - // The iterable is the `right` field — may be identifier, attribute, or call. const rightNode = node.childForFieldName('right'); let iterableName: string | undefined; let methodName: string | undefined; let callExprElementType: string | undefined; + let isEnumerate = false; + + // Extract iterable info from the `right` field — may be identifier, attribute, or call. if (rightNode?.type === 'identifier') { iterableName = rightNode.text; } else if (rightNode?.type === 'attribute') { const prop = rightNode.lastNamedChild; if (prop) iterableName = prop.text; } else if (rightNode?.type === 'call') { - // data.items() → call > function: attribute > identifier('data') + identifier('items') - // get_users() → call > function: identifier (Phase 7.3 — return-type path) const fn = rightNode.childForFieldName('function'); - if (fn?.type === 'attribute') { - const obj = fn.firstNamedChild; - if (obj?.type === 'identifier') iterableName = obj.text; - // Extract method name: items, keys, values - const method = fn.lastNamedChild; - if (method?.type === 'identifier' && method !== obj) methodName = method.text; + if (fn?.type === 'identifier' && fn.text === 'enumerate') { + // enumerate(iterable) or enumerate(d.items()) — unwrap to inner iterable. + isEnumerate = true; + const innerArg = rightNode.childForFieldName('arguments')?.firstNamedChild; + if (innerArg?.type === 'identifier') { + iterableName = innerArg.text; + } else if (innerArg?.type === 'call') { + const extracted = extractMethodCall(innerArg); + if (extracted) ({ iterableName, methodName } = extracted); + } + } else if (fn?.type === 'attribute') { + // data.items() → call > function: attribute > identifier('data') + identifier('items') + const extracted = extractMethodCall(rightNode); + if (extracted) ({ iterableName, methodName } = extracted); } else if (fn?.type === 'identifier') { - // Direct function call: for user in get_users() + // Direct function call: for user in get_users() (Phase 7.3 — return-type path) const rawReturn = returnTypeLookup.lookupRawReturnType(fn.text); if (rawReturn) callExprElementType = extractElementTypeFromString(rawReturn); } @@ -278,11 +320,12 @@ const extractForLoopBinding: ForLoopExtractor = (node, { scopeEnv, declarationTy const leftNode = node.childForFieldName('left'); if (!leftNode) return; - // Handle tuple unpacking: for key, value in data.items() - if (leftNode.type === 'pattern_list') { - const lastChild = leftNode.lastNamedChild; - if (lastChild?.type === 'identifier') { - scopeEnv.set(lastChild.text, elementType); + if (leftNode.type === 'pattern_list' || leftNode.type === 'tuple_pattern') { + // Tuple unpacking: `key, value` or `i, (k, v)` or `(k, v)` — bind the last identifier to element type. + // With enumerate, skip binding if there's only one var (just the index, no value to bind). + const vars = collectPatternIdentifiers(leftNode); + if (vars.length > 0 && (!isEnumerate || vars.length > 1)) { + scopeEnv.set(vars[vars.length - 1].text, elementType); } return; } diff --git a/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/app.py b/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/app.py new file mode 100644 index 0000000000..be3a8ff56e --- /dev/null +++ b/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/app.py @@ -0,0 +1,20 @@ +from user import User +from typing import List + + +def process_users(users: dict[str, User]): + # 3-variable enumerate: i=index, k=key, v=value (User) + for i, k, v in enumerate(users.items()): + v.save() + + +def process_nested_tuple(users: dict[str, User]): + # Nested tuple pattern: i=index, (k,v) tuple unpacked + for i, (k, v) in enumerate(users.items()): + v.save() + + +def process_parenthesized_tuple(users: List[User]): + # Parenthesized tuple as top-level pattern + for (i, u) in enumerate(users): + u.save() diff --git a/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/repo.py b/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/repo.py new file mode 100644 index 0000000000..ec93e0e2f4 --- /dev/null +++ b/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/repo.py @@ -0,0 +1,3 @@ +class Repo: + def save(self) -> bool: + return True diff --git a/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/user.py b/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/user.py new file mode 100644 index 0000000000..c0ddac39f8 --- /dev/null +++ b/gitnexus/test/fixtures/lang-resolution/python-enumerate-loop/user.py @@ -0,0 +1,6 @@ +class User: + def __init__(self, name: str): + self.name = name + + def save(self) -> bool: + return True diff --git a/gitnexus/test/integration/resolvers/python.test.ts b/gitnexus/test/integration/resolvers/python.test.ts index 572f86ea91..c31531124b 100644 --- a/gitnexus/test/integration/resolvers/python.test.ts +++ b/gitnexus/test/integration/resolvers/python.test.ts @@ -1224,3 +1224,60 @@ describe('Python for-loop call_expression iterable resolution (Phase 7.3)', () = expect(saveCalls.length).toBe(1); }); }); + +// --------------------------------------------------------------------------- +// enumerate() for-loop: for i, k, v in enumerate(d.items()) +// --------------------------------------------------------------------------- + +describe('Python enumerate() for-loop resolution', () => { + let result: PipelineResult; + + beforeAll(async () => { + result = await runPipelineFromRepo( + path.join(FIXTURES, 'python-enumerate-loop'), + () => {}, + ); + }, 60000); + + it('detects User class with save method', () => { + expect(getNodesByLabel(result, 'Class')).toContain('User'); + }); + + it('resolves v.save() in enumerate(users.items()) loop to User#save', () => { + // for i, k, v in enumerate(users.items()): v.save() + // v must bind to User (value type of dict[str, User]). + // Without enumerate() support, v is unbound → resolver emits 0 CALLS. + const calls = getRelationships(result, 'CALLS'); + const userSave = calls.find(c => + c.target === 'save' && c.source === 'process_users' && c.targetFilePath?.includes('user.py'), + ); + expect(userSave).toBeDefined(); + }); + + it('does NOT resolve v.save() to a non-User target', () => { + // i is the int index from enumerate — must not produce a spurious CALLS edge + const calls = getRelationships(result, 'CALLS'); + const wrongSave = calls.find(c => + c.target === 'save' && c.source === 'process_users' && !c.targetFilePath?.includes('user.py'), + ); + expect(wrongSave).toBeUndefined(); + }); + + it('resolves nested tuple pattern: for i, (k, v) in enumerate(d.items())', () => { + // Nested tuple_pattern inside pattern_list — must descend to find v + const calls = getRelationships(result, 'CALLS'); + const userSave = calls.find(c => + c.target === 'save' && c.source === 'process_nested_tuple' && c.targetFilePath?.includes('user.py'), + ); + expect(userSave).toBeDefined(); + }); + + it('resolves parenthesized tuple: for (i, u) in enumerate(users)', () => { + // tuple_pattern as top-level left node (not pattern_list) + const calls = getRelationships(result, 'CALLS'); + const userSave = calls.find(c => + c.target === 'save' && c.source === 'process_parenthesized_tuple' && c.targetFilePath?.includes('user.py'), + ); + expect(userSave).toBeDefined(); + }); +}); diff --git a/gitnexus/test/unit/type-env.test.ts b/gitnexus/test/unit/type-env.test.ts index aef07077b0..e740818ba0 100644 --- a/gitnexus/test/unit/type-env.test.ts +++ b/gitnexus/test/unit/type-env.test.ts @@ -2956,6 +2956,50 @@ def process(data: dict[str, User]): expect(flatGet(env, 'user')).toBe('User'); }); + it('Python enumerate(dict.items()): for i, k, v — skips int index, binds value var to User', () => { + // enumerate() wraps the iterable: right node is call with fn='enumerate', not fn.attribute. + // Without enumerate() support, iterableName is never set → v stays unbound. + // With the fix: unwrap enumerate → inner call → data.items() → v binds to User. + const tree = parse(` +def process(data: dict[str, User]): + for i, k, v in enumerate(data.items()): + v.save() + `, Python); + const { env } = buildTypeEnv(tree, 'python'); + expect(flatGet(env, 'v')).toBe('User'); + // i is the int index from enumerate — must NOT be bound to User + expect(flatGet(env, 'i')).toBeUndefined(); + }); + + it('Python enumerate(dict.items()) with nested tuple: for i, (k, v) — binds v to User', () => { + // Nested tuple pattern: `(k, v)` is a tuple_pattern inside the pattern_list. + // AST: pattern_list > [identifier('i'), tuple_pattern > [identifier('k'), identifier('v')]] + // Must descend into tuple_pattern to extract v, not just collect top-level identifiers. + const tree = parse(` +def process(data: dict[str, User]): + for i, (k, v) in enumerate(data.items()): + v.save() + `, Python); + const { env } = buildTypeEnv(tree, 'python'); + expect(flatGet(env, 'v')).toBe('User'); + expect(flatGet(env, 'i')).toBeUndefined(); + }); + + it('Python enumerate with parenthesized tuple: for (k, v) in enumerate(users) — binds v to User', () => { + // Parenthesized tuple pattern: `(k, v)` is a tuple_pattern, not pattern_list. + // AST: for_statement > left: tuple_pattern > [identifier('k'), identifier('v')] + // Must handle tuple_pattern as top-level left node, not just nested inside pattern_list. + const tree = parse(` +def process(users: List[User]): + for (k, v) in enumerate(users): + v.save() + `, Python); + const { env } = buildTypeEnv(tree, 'python'); + // enumerate yields (index, element) — k is int (unbound), v is User + expect(flatGet(env, 'v')).toBe('User'); + expect(flatGet(env, 'k')).toBeUndefined(); + }); + it('TS instanceof narrowing: if (x instanceof User) — first-writer-wins, not block-scoped', () => { // Binds x to User via extractPatternBinding on binary_expression. // Only works when x has no prior type binding in scopeEnv.