From bd993b12fa37838795c70a26f6e7dd6ce0f02855 Mon Sep 17 00:00:00 2001 From: Colby Mchenry Date: Sun, 27 Sep 2026 12:32:27 +0000 Subject: [PATCH] fix(extraction): index Python body docstrings across backends (#1905) (#2004) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(extraction): index Python docstrings, not just preceding comments (#1905) getPrecedingDocstring only walks preceding comment siblings, so a Python docstring — a bare string literal first in the body — never reached the docstring column, never entered nodes_fts, and was never shown by `codegraph node`. The identical sentence written as a leading `#` comment was both. For a Python codebase that is most of the prose there is. (#1905) Adds an optional getBodyDocstring() to LanguageExtractor, in the same shape as getSignature(), and routes every docstring call site through one docstringFor() helper that consults both sources. A node carrying a comment AND a docstring keeps both, joined: they are two things the author wrote about the same symbol and the column is free text. The Python implementation reads the grammar`s string_content rather than slicing quotes off the raw text, so r/u/b prefixes and both triple-quote forms work without a regex per case; f-strings are skipped because an interpolated string is code, not prose. Dedent follows PEP 257 (first line exempt). Scoped to Python deliberately. Julia (a string sibling before the def) and Elixir (@doc) fit the same hook and are left as follow-ups. * fix(extraction): index Python body docstrings across backends (#1905) Python extraction only consulted preceding comments, leaving body docstrings absent from search and rendered prose. Extend the contributor's hook to module and definition docstrings, and mirror it in the native kernel. Handle comments, concatenated literals, and indentation consistently while rejecting bytes, f-strings, and tuples. Verify persistence, exact-name ranking, MCP/CLI rendering, and native/WASM parity. Co-authored-by: Max Hsu Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Max Hsu Co-authored-by: Claude Opus 5.5 --- .../fixtures/kernel-parity/docstrings.py | 102 ++++++++++ __tests__/kernel-tsjs-parity.test.ts | 7 + __tests__/python-body-docstrings.test.ts | 181 ++++++++++++++++++ codegraph-kernel/src/python.rs | 97 +++++++++- src/extraction/languages/python.ts | 62 ++++++ src/extraction/tree-sitter-types.ts | 7 + src/extraction/tree-sitter.ts | 39 ++-- 7 files changed, 480 insertions(+), 15 deletions(-) create mode 100644 __tests__/fixtures/kernel-parity/docstrings.py create mode 100644 __tests__/python-body-docstrings.test.ts diff --git a/__tests__/fixtures/kernel-parity/docstrings.py b/__tests__/fixtures/kernel-parity/docstrings.py new file mode 100644 index 000000000..127fb9222 --- /dev/null +++ b/__tests__/fixtures/kernel-parity/docstrings.py @@ -0,0 +1,102 @@ +# Module header is not the module docstring. +(u"Ledger module " r"documentation.") + +# Legacy reconciliation path. +@decorator +async def reconcile_ledger(): + # A comment before the docstring is allowed. + """Settle the nightly discrepancy with the bank.""" + return 1 + +# Settle the nightly discrepancy with the bank. +def audit_ledger(): + return 1 + +# Ledger class comment. +@decorator +class Ledger: + # Class body comment. + '''Post entries. + + Preserve relative indentation: + details here. + ''' + + # Method comment. + @staticmethod + async def settle(): + # Method body comment. + (U"Settle " # Between concatenated strings. + R"the ledger.") + return 1 + +def raw_doc(): + R'''Raw \d+ expression.''' + return 1 + +def empty_concat(): + "" "Kept prose." "" + return 1 + +def empty_doc(): + "" + return 1 + +def bytes_doc(): + b"Not documentation." + return 1 + +def raw_bytes_doc(): + BR"Not documentation." + return 1 + +def formatted_doc(): + F"Not documentation." + return 1 + +def interpolated_doc(): + rf"Not {documentation}." + return 1 + +def concatenated_bytes(): + b"Not " b"documentation." + return 1 + +def concatenated_format(): + "Not " f"documentation." + return 1 + +def late_string(): + value = 1 + "Not documentation." + return value + +def computed_string(): + "Not " + "documentation." + return 1 + +def tuple_string(): + ("Not documentation.",) + return 1 + +def bare_tuple(): + "Not documentation.", "Also not." + return 1 + +def singleton_tuple(): + "Not documentation.", + return 1 + +# +def blank_comment(): + (("Kept despite blank comment.")) + return 1 + +def nested_string(): + if True: + "Not documentation." + return 1 + +def ranked_decoy(): + """reconcile_ledger reconcile_ledger reconcile_ledger""" + return 1 diff --git a/__tests__/kernel-tsjs-parity.test.ts b/__tests__/kernel-tsjs-parity.test.ts index 6727ca93f..3fe2834ab 100644 --- a/__tests__/kernel-tsjs-parity.test.ts +++ b/__tests__/kernel-tsjs-parity.test.ts @@ -144,6 +144,13 @@ function nested(holder) { assertParity('fixtures/torture.py', fs.readFileSync(file, 'utf8'), 'python'); }); + it.each(['LF', 'CRLF'])('Python body docstrings parity (%s, #1905)', (ending) => { + const source = fs.readFileSync(path.join(FIXTURE_DIR, 'docstrings.py'), 'utf8'); + const result = assertParity('ledger.py', ending === 'CRLF' ? source.replace(/\n/g, '\r\n') : source, 'python'); + expect(result.nodes.find((n) => n.kind === 'file')?.docstring).toBe('Ledger module documentation.'); + expect(result.nodes.find((n) => n.name === 'settle')?.docstring).toBe('Method comment.\n\nSettle the ledger.'); + }); + it('torture fixture (go): receivers, embedding, interfaces, composite literals', () => { const file = path.join(FIXTURE_DIR, 'torture.go'); assertParity('fixtures/torture.go', fs.readFileSync(file, 'utf8'), 'go'); diff --git a/__tests__/python-body-docstrings.test.ts b/__tests__/python-body-docstrings.test.ts new file mode 100644 index 000000000..909805ca3 --- /dev/null +++ b/__tests__/python-body-docstrings.test.ts @@ -0,0 +1,181 @@ +/** + * Python body docstrings (#1905) + * + * `getPrecedingDocstring` only walks preceding comment siblings, so a Python + * docstring — a bare string literal first in the body — never reached the + * `docstring` column, never entered `nodes_fts`, and was never shown. The + * identical sentence written as a leading `#` comment was both. + */ + +import { describe, it, expect, beforeAll, beforeEach, afterEach, vi } from 'vitest'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import { execFileSync } from 'child_process'; +import { CodeGraph } from '../src'; +import { ToolHandler } from '../src/mcp/tools'; +import { getKernel, resetKernelForTests } from '../src/extraction/kernel'; +import { extractFromSource } from '../src/extraction'; +import { initGrammars, loadAllGrammars } from '../src/extraction/grammars'; + +beforeAll(async () => { + await initGrammars(); + await loadAllGrammars(); +}); + +function docstrings(code: string): Map { + return new Map(extractFromSource('ledger.py', code).nodes.map((n) => [n.name, n.docstring])); +} + +const fixture = fs.readFileSync(path.join(__dirname, 'fixtures/kernel-parity/docstrings.py'), 'utf8'); +const kernelBuilt = fs.existsSync(path.join(__dirname, '../codegraph-kernel/prebuilds', + `${process.platform}-${process.arch}`, 'codegraph-kernel.node')); + +describe.each(['native', 'wasm'].filter((backend) => backend === 'wasm' || kernelBuilt))( + 'Python body docstrings (%s)', (backend) => { + let dir: string | undefined; + let cg: CodeGraph | undefined; + + beforeEach(() => { + vi.stubEnv('CODEGRAPH_KERNEL', backend === 'wasm' ? '0' : '1'); + vi.stubEnv('CODEGRAPH_KERNEL_LANGS', 'all'); + resetKernelForTests(); + if (backend === 'native') expect(getKernel()).not.toBeNull(); + }); + + afterEach(() => { + cg?.close(); + cg = undefined; + if (dir) fs.rmSync(dir, { recursive: true, force: true }); + dir = undefined; + vi.unstubAllEnvs(); + resetKernelForTests(); + }); + + it.each(['LF', 'CRLF'])('handles module, decorated/async definitions and literal forms (%s)', (ending) => { + const byName = docstrings(ending === 'CRLF' ? fixture.replace(/\n/g, '\r\n') : fixture); + expect(byName.get('ledger.py')).toBe('Ledger module documentation.'); + expect(byName.get('reconcile_ledger')).toBe('Legacy reconciliation path.\n\nSettle the nightly discrepancy with the bank.'); + expect(byName.get('Ledger')).toBe('Ledger class comment.\n\nPost entries.\n\nPreserve relative indentation:\n details here.'); + expect(byName.get('settle')).toBe('Method comment.\n\nSettle the ledger.'); + expect(byName.get('raw_doc')).toBe('Raw \\d+ expression.'); + expect(byName.get('empty_concat')).toBe('Kept prose.'); + expect(byName.get('blank_comment')).toBe('Kept despite blank comment.'); + for (const name of ['empty_doc', 'bytes_doc', 'raw_bytes_doc', 'formatted_doc', + 'interpolated_doc', 'concatenated_bytes', 'concatenated_format', 'late_string', + 'computed_string', 'tuple_string', 'bare_tuple', 'singleton_tuple', 'nested_string']) { + expect(byName.has(name)).toBe(true); + expect(byName.get(name), name).toBeUndefined(); + } + }); + + it('expands tabs and preserves relative indentation', () => { + expect(docstrings('def tabbed():\n """Summary.\n\n\tDetails.\n\t Nested.\n """\n').get('tabbed')) + .toBe('Summary.\n\nDetails.\n Nested.'); + }); + + it.each(['b"bytes"', 'f"formatted"', 'pass\n"late"', '"one" + "two"'])( + 'rejects a non-docstring module first statement: %s', (source) => { + expect(docstrings(source).get('ledger.py')).toBeUndefined(); + }, + ); + + it('persists prose for FTS, keeps exact names first, and renders it through MCP and CLI', async () => { + dir = fs.mkdtempSync(path.join(os.tmpdir(), 'cg-1905-')); + fs.writeFileSync(path.join(dir, 'ledger.py'), fixture); + cg = CodeGraph.initSync(dir); + await cg.indexAll(); + cg.close(); + cg = CodeGraph.openSync(dir); + expect(cg.getNodesByName('reconcile_ledger')[0].docstring) + .toBe('Legacy reconciliation path.\n\nSettle the nightly discrepancy with the bank.'); + expect(cg.searchNodes('discrepancy bank').map(({ node }) => node.name)) + .toEqual(expect.arrayContaining(['reconcile_ledger', 'audit_ledger'])); + expect(cg.searchNodes('reconcile_ledger')[0].node.name).toBe('reconcile_ledger'); + expect(cg.searchNodes('documentation').some(({ node }) => node.kind === 'file' && node.name === 'ledger.py')).toBe(true); + const result = await new ToolHandler(cg).execute('codegraph_node', { symbol: 'reconcile_ledger', includeCode: true }); + expect(result.isError).not.toBe(true); + const text = result.content.filter((c) => c.type === 'text').map((c) => c.text).join('\n'); + // Prose must occur before the verbatim source, not only inside it. + expect(text.split('```')[0]).toContain('Settle the nightly discrepancy with the bank.'); + const cli = execFileSync(process.execPath, [path.join(__dirname, '../dist/bin/codegraph.js'), + 'node', 'reconcile_ledger', '--path', dir], { encoding: 'utf8', env: process.env }); + expect(cli.split('```')[0]).toContain('Settle the nightly discrepancy with the bank.'); + }); + + it('extracts the same sentence from a docstring as from a leading comment (#1905)', () => { + // The issue's repro: one sentence, two positions. Before this change only + // `audit_ledger` carried it. + const byName = docstrings(` +def reconcile_ledger(): + """Settle the nightly discrepancy with the bank.""" + return LEDGER + +# Settle the nightly discrepancy with the bank. +def audit_ledger(): + return LEDGER +`); + expect(byName.get('reconcile_ledger')).toBe('Settle the nightly discrepancy with the bank.'); + expect(byName.get('audit_ledger')).toBe('Settle the nightly discrepancy with the bank.'); + }); + + it('dedents a multi-line docstring to its own left margin', () => { + const byName = docstrings(` +class Ledger: + """Post entries to the general ledger. + + Nightly reconciliation runs against the bank feed. + """ + pass +`); + expect(byName.get('Ledger')).toBe( + 'Post entries to the general ledger.\n\nNightly reconciliation runs against the bank feed.', + ); + }); + + it("keeps both when a definition carries a comment AND a docstring", () => { + // Two things the author wrote about the same symbol; the column is free + // text, so neither is dropped. + const byName = docstrings(` +# Legacy path, kept for the 2019 import. +def reconcile_ledger(): + """Settle the nightly discrepancy with the bank.""" + return LEDGER +`); + expect(byName.get('reconcile_ledger')).toBe( + 'Legacy path, kept for the 2019 import.\n\nSettle the nightly discrepancy with the bank.', + ); + }); + + it('handles the single-quote and prefixed forms', () => { + const byName = docstrings(` +def with_single(): + '''Single-quoted docstring.''' + return 1 + +def with_raw_prefix(): + r"""Raw \\d+ docstring.""" + return 2 +`); + expect(byName.get('with_single')).toBe('Single-quoted docstring.'); + expect(byName.get('with_raw_prefix')).toBe('Raw \\d+ docstring.'); + }); + + it('skips an f-string first statement — that is code, not prose', () => { + const byName = docstrings(` +def interpolated(name): + f"""Hello {name}.""" + return name +`); + expect(byName.get('interpolated') ?? null).toBeNull(); + }); + + it('does not treat a non-string first statement as prose', () => { + const byName = docstrings(` +def no_docstring(): + total = 0 + return total +`); + expect(byName.get('no_docstring') ?? null).toBeNull(); + }); +}); diff --git a/codegraph-kernel/src/python.rs b/codegraph-kernel/src/python.rs index 79a825c1c..66f61f885 100644 --- a/codegraph-kernel/src/python.rs +++ b/codegraph-kernel/src/python.rs @@ -21,6 +21,93 @@ use tree_sitter::{Node, Parser}; const MAX_VALUE_REF_NODES: usize = 20_000; +// Keep body-docstring selection and cleaning in sync with languages/python.ts. +fn body_docstring(node: Node, src: &str) -> Option { + let body = if node.kind() == "module" { node } else { node.child_by_field_name("body")? }; + let first = (0..body.named_child_count()) + .filter_map(|i| body.named_child(i)) + .find(|c| c.kind() != "comment")?; + if first.kind() != "expression_statement" { + return None; + } + if first.named_child_count() != 1 || (0..first.child_count()) + .filter_map(|i| first.child(i)).any(|c| c.kind() == ",") { + return None; + } + let mut literal = first.named_child(0)?; + while literal.kind() == "parenthesized_expression" { + literal = (0..literal.named_child_count()) + .filter_map(|i| literal.named_child(i)) + .find(|c| c.kind() != "comment")?; + } + let strings = if literal.kind() == "concatenated_string" { + (0..literal.named_child_count()).filter_map(|i| literal.named_child(i)) + .filter(|c| c.kind() != "comment").collect::>() + } else { + vec![literal] + }; + let mut raw = String::new(); + for string in strings { + if string.kind() != "string" { + return None; + } + let start = (0..string.named_child_count()).filter_map(|i| string.named_child(i)) + .find(|c| c.kind() == "string_start")?; + if src[start.byte_range()].bytes().any(|b| matches!(b, b'b' | b'B' | b'f' | b'F')) { + return None; + } + if !(0..string.named_child_count()).filter_map(|i| string.named_child(i)) + .any(|c| c.kind() == "string_end") { + return None; + } + for i in 0..string.named_child_count() { + if let Some(content) = string.named_child(i).filter(|c| c.kind() == "string_content") { + raw.push_str(&src[content.byte_range()]); + } + } + } + let cleaned = dedent_docstring(&raw); + if cleaned.is_empty() { None } else { Some(cleaned) } +} + +fn docstring_space(c: char) -> bool { + // Match JavaScript trim, including BOM but excluding NEXT LINE. + (c.is_whitespace() && c != '\u{85}') || c == '\u{feff}' +} + +fn dedent_docstring(raw: &str) -> String { + let normalized = raw.replace("\r\n", "\n").replace('\r', "\n"); + let lines: Vec = normalized.split('\n').map(|line| { + let mut column = 0; + let mut expanded = String::new(); + for c in line.chars() { + let width = if c == '\t' { 8 - column % 8 } else { 1 }; + column += width; + if c == '\t' { expanded.push_str(&" ".repeat(width)); } else { expanded.push(c); } + } + expanded + }).collect(); + let indent = lines.iter().skip(1).filter(|l| !l.trim_matches(docstring_space).is_empty()) + .map(|l| l.chars().take_while(|c| docstring_space(*c)).count()).min().unwrap_or(0); + let out: Vec = lines.iter().enumerate().map(|(i, l)| { + if i == 0 { l.trim_matches(docstring_space).to_string() } else { + l.chars().skip(indent).collect::().trim_end_matches(docstring_space).to_string() + } + }).collect(); + let start = out.iter().position(|l| !l.trim_matches(docstring_space).is_empty()).unwrap_or(out.len()); + let end = out.iter().rposition(|l| !l.trim_matches(docstring_space).is_empty()).map_or(start, |i| i + 1); + out[start..end].join("\n") +} + +fn docstring_for(node: Node, src: &str) -> Option { + let preceding = preceding_docstring(node, src); + let body = body_docstring(node, src); + match (preceding, body) { + (Some(a), Some(b)) if !a.is_empty() => Some(format!("{a}\n\n{b}")), + (a, b) => b.or(a), + } +} + struct Scope { row: u32, kind: &'static str, @@ -102,6 +189,8 @@ pub fn extract(file_path: &str, source: &str) -> Result { let file_id = w.arena.put(&ids::file_node_id(file_path)); let name_ref = w.arena.put(base_name); let qn_ref = w.arena.put(file_path); + let docstring = body_docstring(tree.root_node(), source) + .map(|doc| w.arena.put(&doc)).unwrap_or(NONE_STR); w.tables.push_node(&NodeRow { kind: node_kind_index("file").unwrap(), visibility: 0, @@ -113,7 +202,7 @@ pub fn extract(file_path: &str, source: &str) -> Result { name: name_ref, qualified_name: qn_ref, id: file_id, - docstring: NONE_STR, + docstring, signature: NONE_STR, decorators: NONE_STR, type_parameters: NONE_STR, @@ -402,7 +491,7 @@ impl<'t> Walker<'t> { return; } let extra = Extra { - docstring: preceding_docstring(node, self.src), + docstring: docstring_for(node, self.src), signature: self.signature_of(node), is_async: Some(self.is_async(node)), is_static: Some(self.is_static(node)), @@ -421,7 +510,7 @@ impl<'t> Walker<'t> { stack_guard!(); let name = self.extract_name(node); let extra = Extra { - docstring: preceding_docstring(node, self.src), + docstring: docstring_for(node, self.src), signature: self.signature_of(node), is_async: Some(self.is_async(node)), is_static: Some(self.is_static(node)), @@ -439,7 +528,7 @@ impl<'t> Walker<'t> { stack_guard!(); let name = self.extract_name(node); let extra = Extra { - docstring: preceding_docstring(node, self.src), + docstring: docstring_for(node, self.src), ..Extra::default() }; let Some(row) = self.create_node("class", &name, node, extra) else { return }; diff --git a/src/extraction/languages/python.ts b/src/extraction/languages/python.ts index 77807d667..99d0c48ac 100644 --- a/src/extraction/languages/python.ts +++ b/src/extraction/languages/python.ts @@ -16,6 +16,40 @@ export const pythonExtractor: LanguageExtractor = { bodyField: 'body', paramsField: 'parameters', returnField: 'return_type', + /** + * Python states intent in a docstring — a bare string literal as the first + * statement of the body — not in a preceding comment, so the comment-sibling + * walk never reached it and the prose never entered the index (#1905). + * + * Reads `string_content` rather than slicing quotes off the raw text: the + * grammar already separates delimiters from body, which keeps `r`/`u` + * prefixes and both `"""` and `'''` forms working without a regex per case. + * Bytes and f-strings are not Python docstrings. + */ + getBodyDocstring: (node, source) => { + const body = node.type === 'module' ? node : getChildByField(node, 'body'); + if (!body) return undefined; + const first = body.namedChildren.find((c) => c.type !== 'comment'); + if (!first || first.type !== 'expression_statement') return undefined; + if (first.namedChildCount !== 1 || first.children.some((c) => c.type === ',')) return undefined; + let literal = first.namedChild(0); + while (literal?.type === 'parenthesized_expression') { + literal = literal.namedChildren.find((c) => c.type !== 'comment') ?? null; + } + if (!literal) return undefined; + const strings = literal.type === 'concatenated_string' + ? literal.namedChildren.filter((c) => c.type !== 'comment') : [literal]; + let raw = ''; + for (const string of strings) { + if (string.type !== 'string') return undefined; + const start = string.namedChildren.find((c) => c.type === 'string_start'); + if (!start || /[bf]/i.test(getNodeText(start, source))) return undefined; + if (!string.namedChildren.some((c) => c.type === 'string_end')) return undefined; + const content = string.namedChildren.find((c) => c.type === 'string_content'); + if (content) raw += getNodeText(content, source); + } + return dedentDocstring(raw) || undefined; + }, getSignature: (node, source) => { const params = getChildByField(node, 'parameters'); const returnType = getChildByField(node, 'return_type'); @@ -51,3 +85,31 @@ export const pythonExtractor: LanguageExtractor = { return null; }, }; + +/** + * A docstring is indented to its definition, so every line after the first + * carries that indentation. Strip the common prefix (PEP 257's rule: the first + * line is exempt because it starts right after the opening quotes) and drop + * blank edges, so the stored prose reads the same as a comment-derived one. + */ +function dedentDocstring(raw: string): string { + const lines = raw.replace(/\r\n?/g, '\n').split('\n').map((line) => { + let column = 0; + return Array.from(line, (char) => { + const width = char === '\t' ? 8 - column % 8 : 1; + column += width; + return char === '\t' ? ' '.repeat(width) : char; + }).join(''); + }); + const rest = lines.slice(1).filter((l) => l.trim().length > 0); + const indent = rest.length === 0 + ? 0 + : Math.min(...rest.map((l) => l.length - l.trimStart().length)); + const out = [ + lines[0]?.trim() ?? '', + ...lines.slice(1).map((l) => l.slice(indent).trimEnd()), + ]; + while (out.length > 0 && out[0]!.trim() === '') out.shift(); + while (out.length > 0 && out[out.length - 1]!.trim() === '') out.pop(); + return out.join('\n'); +} diff --git a/src/extraction/tree-sitter-types.ts b/src/extraction/tree-sitter-types.ts index 19f749cf5..e344d545a 100644 --- a/src/extraction/tree-sitter-types.ts +++ b/src/extraction/tree-sitter-types.ts @@ -152,6 +152,13 @@ export interface LanguageExtractor { /** Extract signature from node */ getSignature?: (node: SyntaxNode, source: string) => string | undefined; + /** + * Extract a docstring that lives INSIDE the node's body rather than as a + * preceding comment sibling — the shape `getPrecedingDocstring` cannot see. + * Python puts a bare string literal first in the body; Julia and Elixir have + * their own variants. Returns cleaned prose, or undefined when there is none. + */ + getBodyDocstring?: (node: SyntaxNode, source: string) => string | undefined; /** Extract visibility from node */ getVisibility?: (node: SyntaxNode) => 'public' | 'private' | 'protected' | 'internal' | undefined; /** Check if node is exported */ diff --git a/src/extraction/tree-sitter.ts b/src/extraction/tree-sitter.ts index d8fed3e5c..dd36407f9 100644 --- a/src/extraction/tree-sitter.ts +++ b/src/extraction/tree-sitter.ts @@ -432,6 +432,21 @@ function isUnresolvedTsJsChain(node: SyntaxNode, source: string): boolean { const REACT_HANDLER_HOOKS = /^(?:React\.)?use(?:Callback|EffectEvent|Event)$/; export class TreeSitterExtractor { + /** + * The node's prose, from either place it can live: a preceding comment + * sibling (every language) or a docstring inside the body (Python's bare + * first-statement string, and the same shape in other languages that opt in + * via `getBodyDocstring`). When a node carries both, they are joined rather + * than one winning — they are two separate things the author wrote about the + * same symbol, and the column holds free text (#1905). + */ + private docstringFor(node: SyntaxNode): string | undefined { + const preceding = getPrecedingDocstring(node, this.source); + const body = this.extractor?.getBodyDocstring?.(node, this.source); + if (preceding && body) return `${preceding}\n\n${body}`; + return body || preceding; + } + private filePath: string; private language: Language; private source: string; @@ -563,6 +578,8 @@ export class TreeSitterExtractor { isExported: false, updatedAt: Date.now(), }; + const fileDocstring = this.extractor?.getBodyDocstring?.(this.tree.rootNode, this.source); + if (fileDocstring) fileNode.docstring = fileDocstring; this.nodes.push(fileNode); // Push file node onto stack so top-level declarations get contains edges @@ -1660,7 +1677,7 @@ export class TreeSitterExtractor { return; } - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const signature = this.extractor.getSignature?.(node, this.source); const visibility = this.extractor.getVisibility?.(node); const isExported = commonJsExport || this.extractor.isExported?.(node, this.source); @@ -1778,7 +1795,7 @@ export class TreeSitterExtractor { if (this.extractor.skipBodilessClass && !resolvedBody) return; const name = extractName(node, this.source, this.extractor); - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const visibility = this.extractor.getVisibility?.(node); const isExported = this.extractor.isExported?.(node, this.source); @@ -1862,7 +1879,7 @@ export class TreeSitterExtractor { return; } - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const signature = this.extractor.getSignature?.(node, this.source); const visibility = this.extractor.getVisibility?.(node); const isAsync = this.extractor.isAsync?.(node); @@ -1928,7 +1945,7 @@ export class TreeSitterExtractor { if (!this.extractor) return; const name = extractName(node, this.source, this.extractor); - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const isExported = this.extractor.isExported?.(node, this.source); const kind: NodeKind = this.extractor.interfaceKind ?? 'interface'; @@ -1987,7 +2004,7 @@ export class TreeSitterExtractor { return; const name = extractName(node, this.source, this.extractor); - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const visibility = this.extractor.getVisibility?.(node); const isExported = this.extractor.isExported?.(node, this.source); @@ -2031,7 +2048,7 @@ export class TreeSitterExtractor { if (!body) return; const name = extractName(node, this.source, this.extractor); - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const visibility = this.extractor.getVisibility?.(node); const isExported = this.extractor.isExported?.(node, this.source); @@ -2097,7 +2114,7 @@ export class TreeSitterExtractor { private extractProperty(node: SyntaxNode): Node | null { if (!this.extractor) return null; - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const visibility = this.extractor.getVisibility?.(node); const isStatic = this.extractor.isStatic?.(node) ?? false; @@ -2167,7 +2184,7 @@ export class TreeSitterExtractor { private extractField(node: SyntaxNode): void { if (!this.extractor) return; - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const visibility = this.extractor.getVisibility?.(node); const isStatic = this.extractor.isStatic?.(node) ?? false; @@ -2704,7 +2721,7 @@ export class TreeSitterExtractor { const isConst = this.extractor.isConst?.(node) ?? false; const kind: NodeKind = isConst ? 'constant' : 'variable'; - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const isExported = this.extractor.isExported?.(node, this.source) ?? false; // Extract variable declarators based on language @@ -3177,7 +3194,7 @@ export class TreeSitterExtractor { const name = extractName(node, this.source, this.extractor); if (name === '') return false; - const docstring = getPrecedingDocstring(node, this.source); + const docstring = this.docstringFor(node); const isExported = this.extractor.isExported?.(node, this.source); // Check if this type alias is actually a struct or interface definition @@ -3341,7 +3358,7 @@ export class TreeSitterExtractor { ? 'method' : this.isTsFunctionTypedProperty(child) ? 'method' : 'property'; - const docstring = getPrecedingDocstring(child, this.source); + const docstring = this.docstringFor(child); const signature = getNodeText(child, this.source); this.createNode(memberKind, memberName, child, { docstring,