From fc4b48d1f11007d168dad1ce8620f7e44aa41617 Mon Sep 17 00:00:00 2001 From: Gergo Magyar Date: Fri, 25 Sep 2026 15:44:45 +0000 Subject: [PATCH] fix(review): locate notebook sources independent of JSON key order Walk the cells array as objects so source-before-cell_type notebooks map correctly, remap assignment lines, and cap the extract cache. Co-authored-by: Cursor --- .../src/core/ingestion/ipynb-extractor.ts | 133 +++++++++++++++--- .../core/ingestion/workers/parse-worker.ts | 2 +- gitnexus/test/unit/ipynb-extractor.test.ts | 19 +++ 3 files changed, 135 insertions(+), 19 deletions(-) diff --git a/gitnexus/src/core/ingestion/ipynb-extractor.ts b/gitnexus/src/core/ingestion/ipynb-extractor.ts index 0423b0aef..327187741 100644 --- a/gitnexus/src/core/ingestion/ipynb-extractor.ts +++ b/gitnexus/src/core/ingestion/ipynb-extractor.ts @@ -87,28 +87,119 @@ function jsonValueSpan(content: string, start: number): { start: number; end: nu return null; } +function jsonBraceSpan( + content: string, + start: number, + open: '{' | '[', + close: '}' | ']', +): { start: number; end: number } | null { + const i = skipWs(content, start); + if (i >= content.length || content[i] !== open) return null; + let depth = 1; + let j = i + 1; + let inStr = false; + while (j < content.length && depth > 0) { + const ch = content[j]; + if (inStr) { + if (ch === '\\') j += 2; + else { + if (ch === '"') inStr = false; + j++; + } + continue; + } + if (ch === '"') inStr = true; + else if (ch === open) depth++; + else if (ch === close) depth--; + j++; + } + if (depth !== 0) return null; + return { start: i, end: j }; +} + +function findDepth1Key(content: string, objStart: number, objEnd: number, key: string): number { + let depth = 0; + let inStr = false; + let j = objStart; + while (j < objEnd) { + const ch = content[j]; + if (inStr) { + if (ch === '\\') j += 2; + else { + if (ch === '"') inStr = false; + j++; + } + continue; + } + if (ch === '"') { + if (depth === 1 && content.startsWith(key, j)) return j; + inStr = true; + j++; + continue; + } + if (ch === '{' || ch === '[') depth++; + else if (ch === '}' || ch === ']') depth--; + j++; + } + return -1; +} + +function findCellsArraySpan(content: string): { start: number; end: number } | null { + const root = jsonBraceSpan(content, 0, '{', '}'); + if (!root) return null; + const cellsKey = findDepth1Key(content, root.start, root.end, '"cells"'); + if (cellsKey < 0) return null; + const colon = content.indexOf(':', cellsKey + 7); + if (colon < 0 || colon >= root.end) return null; + return jsonBraceSpan(content, colon + 1, '[', ']'); +} + function findNextCodeCellSourceSpan( content: string, from: number, ): { span: { start: number; end: number }; nextFrom: number } | null { - let search = from; - while (search < content.length) { - const typeKey = content.indexOf('"cell_type"', search); - if (typeKey < 0) return null; - const colon = content.indexOf(':', typeKey + 11); - if (colon < 0) return null; - const valueStart = skipWs(content, colon + 1); - if (content.slice(valueStart, valueStart + 6) !== '"code"') { - search = typeKey + 11; + const cells = findCellsArraySpan(content); + if (!cells) return null; + let search = Math.max(from, cells.start + 1); + while (search < cells.end) { + const brace = content.indexOf('{', search); + if (brace < 0 || brace >= cells.end) return null; + const obj = jsonBraceSpan(content, brace, '{', '}'); + if (!obj || obj.end > cells.end) { + search = brace + 1; + continue; + } + const typeKey = findDepth1Key(content, obj.start, obj.end, '"cell_type"'); + if (typeKey < 0) { + search = obj.end; + continue; + } + const colon = content.indexOf(':', typeKey + 11); + if (colon < 0 || colon >= obj.end) { + search = obj.end; + continue; + } + const valueStart = skipWs(content, colon + 1); + if (content.slice(valueStart, valueStart + 6) !== '"code"') { + search = obj.end; + continue; + } + const sourceKey = findDepth1Key(content, obj.start, obj.end, '"source"'); + if (sourceKey < 0) { + search = obj.end; continue; } - const sourceKey = content.indexOf('"source"', valueStart); - if (sourceKey < 0) return null; const srcColon = content.indexOf(':', sourceKey + 8); - if (srcColon < 0) return null; + if (srcColon < 0 || srcColon >= obj.end) { + search = obj.end; + continue; + } const span = jsonValueSpan(content, srcColon + 1); - if (!span) return null; - return { span, nextFrom: span.end }; + if (!span) { + search = obj.end; + continue; + } + return { span, nextFrom: obj.end }; } return null; } @@ -231,13 +322,13 @@ export function extractNotebookPython(content: string): NotebookPythonExtraction if (chunks.length > 0) { chunks.push('\n\n'); } + const lineCount = text.split('\n').length; + const extractStartLine = + segments.length === 0 ? 0 : segments[segments.length - 1].extractEndLine + 2; chunks.push(text); - const assembled = chunks.join(''); - const extractStartLine = indexToLine(assembled, assembled.length - text.length); - const extractEndLine = indexToLine(assembled, assembled.length - 1); segments.push({ extractStartLine, - extractEndLine, + extractEndLine: extractStartLine + lineCount - 1, jsonStartLine, jsonEndLine, }); @@ -268,6 +359,8 @@ const extractCache = new Map< >(); /** Memoize extraction for FTS/CSV (same file, many symbols). */ +const EXTRACT_CACHE_LIMIT = 32; + export function extractNotebookPythonCached( filePath: string, content: string, @@ -275,6 +368,10 @@ export function extractNotebookPythonCached( const hit = extractCache.get(filePath); if (hit && hit.content === content) return hit.result; const result = extractNotebookPython(content); + if (extractCache.size >= EXTRACT_CACHE_LIMIT && !extractCache.has(filePath)) { + const oldest = extractCache.keys().next().value; + if (oldest !== undefined) extractCache.delete(oldest); + } extractCache.set(filePath, { content, result }); return result; } diff --git a/gitnexus/src/core/ingestion/workers/parse-worker.ts b/gitnexus/src/core/ingestion/workers/parse-worker.ts index 1d80e4322..4f9dbab36 100644 --- a/gitnexus/src/core/ingestion/workers/parse-worker.ts +++ b/gitnexus/src/core/ingestion/workers/parse-worker.ts @@ -1896,7 +1896,7 @@ const processFileGroup = ( sourceId: srcId, receiverText, propertyName, - line: captureMap['assignment'].startPosition.row + 1, + line: mapRow(captureMap['assignment'].startPosition.row) + 1, ...(receiverTypeName ? { receiverTypeName } : {}), }); } diff --git a/gitnexus/test/unit/ipynb-extractor.test.ts b/gitnexus/test/unit/ipynb-extractor.test.ts index 35bb6d98b..ad10e5c92 100644 --- a/gitnexus/test/unit/ipynb-extractor.test.ts +++ b/gitnexus/test/unit/ipynb-extractor.test.ts @@ -40,6 +40,25 @@ describe('isPythonFamilyLanguage', () => { }); describe('extractNotebookPython', () => { + it('maps source that is serialized before cell_type', () => { + const content = JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + metadata: { kernelspec: { language: 'python', name: 'python3', display_name: 'Python' } }, + cells: [ + { + source: ['def train():\n', ' pass\n'], + cell_type: 'code', + metadata: {}, + outputs: [], + }, + ], + }); + const result = extractNotebookPython(content); + expect(result?.pythonSource).toContain('def train'); + expect(result?.segments).toHaveLength(1); + }); + it('extracts def train from a Python v4 notebook', () => { const content = notebook({ language: 'python',