python.rs 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978
  1. //! Python extraction — a faithful Rust port of `TreeSitterExtractor`'s Python
  2. //! paths (src/extraction/tree-sitter.ts) plus languages/python.ts.
  3. //!
  4. //! Same porting contract as tsjs/java: behavior parity, bug-for-bug —
  5. //! including the quirks: decorates refs only fire for bare-identifier
  6. //! decorators (`@staticmethod` yes, `@app.route(...)` no — python's `call`
  7. //! kind isn't `call_expression`), module-level assignments always extract as
  8. //! `variable` (no isConst hook), and `self.method` fn-ref candidates carry the
  9. //! BARE attribute name. Python is not a TYPE_ANNOTATION language — no type
  10. //! refs anywhere. Files with parse errors defer to wasm.
  11. use crate::buffers::{
  12. build_meta, edge_kind_index, node_kind_index, Arena, BoolFlags, EdgeRow, EmitOut, NodeRow,
  13. RefRow, StrRef, Tables, FLAG_IS_ASYNC, FLAG_IS_STATIC, FUNCTION_REF_CODE, NONE, NONE_STR,
  14. };
  15. use crate::docstring::preceding_docstring;
  16. use crate::ids;
  17. use crate::textutil as util;
  18. use std::collections::{HashMap, HashSet};
  19. use tree_sitter::{Node, Parser};
  20. const MAX_VALUE_REF_NODES: usize = 20_000;
  21. struct Scope {
  22. row: u32,
  23. kind: &'static str,
  24. name: String,
  25. }
  26. #[derive(Default)]
  27. struct Extra {
  28. docstring: Option<String>,
  29. signature: Option<String>,
  30. is_async: Option<bool>,
  31. is_static: Option<bool>,
  32. }
  33. struct ValueScope<'t> {
  34. row: u32,
  35. node: Node<'t>,
  36. name: String,
  37. }
  38. struct Cand {
  39. from: u32,
  40. name: String,
  41. line: u32,
  42. column_byte: usize,
  43. row: usize,
  44. }
  45. pub struct Walker<'t> {
  46. src: &'t str,
  47. file_path: &'t str,
  48. line_starts: Vec<usize>,
  49. arena: Arena,
  50. tables: Tables,
  51. stack: Vec<Scope>,
  52. node_ids: Vec<String>,
  53. defined_fn_names: HashSet<String>,
  54. imported_names: HashSet<String>,
  55. fn_ref_cands: Vec<Cand>,
  56. fs_values: HashMap<String, u32>,
  57. fs_value_counts: HashMap<String, u32>,
  58. value_scopes: Vec<ValueScope<'t>>,
  59. }
  60. pub fn extract(file_path: &str, source: &str) -> Result<EmitOut, String> {
  61. let grammar = crate::langs::grammar_for("python").ok_or("no python grammar")?;
  62. let t0 = std::time::Instant::now();
  63. let mut parser = Parser::new();
  64. parser
  65. .set_language(&grammar)
  66. .map_err(|e| format!("set_language(python) failed: {e}"))?;
  67. let tree = parser
  68. .parse(source, None)
  69. .ok_or_else(|| "parser returned null tree".to_string())?;
  70. if tree.root_node().has_error() {
  71. return Err("defer: parse tree contains errors — wasm recovery is canonical".to_string());
  72. }
  73. let mut w = Walker {
  74. src: source,
  75. file_path,
  76. line_starts: util::line_starts(source),
  77. arena: Arena::default(),
  78. tables: Tables::default(),
  79. stack: Vec::new(),
  80. node_ids: Vec::new(),
  81. defined_fn_names: HashSet::new(),
  82. imported_names: HashSet::new(),
  83. fn_ref_cands: Vec::new(),
  84. fs_values: HashMap::new(),
  85. fs_value_counts: HashMap::new(),
  86. value_scopes: Vec::new(),
  87. };
  88. let line_count = source.bytes().filter(|b| *b == b'\n').count() as u32 + 1;
  89. let base_name = file_path.rsplit(['/', '\\']).next().unwrap_or(file_path);
  90. let mut flags = BoolFlags::default();
  91. flags.set(crate::buffers::FLAG_IS_EXPORTED, false);
  92. let file_id = w.arena.put(&ids::file_node_id(file_path));
  93. let name_ref = w.arena.put(base_name);
  94. let qn_ref = w.arena.put(file_path);
  95. w.tables.push_node(&NodeRow {
  96. kind: node_kind_index("file").unwrap(),
  97. visibility: 0,
  98. flags,
  99. start_line: 1,
  100. end_line: line_count,
  101. start_column: 0,
  102. end_column: 0,
  103. name: name_ref,
  104. qualified_name: qn_ref,
  105. id: file_id,
  106. docstring: NONE_STR,
  107. signature: NONE_STR,
  108. decorators: NONE_STR,
  109. type_parameters: NONE_STR,
  110. return_type: NONE_STR,
  111. extra_json: NONE_STR,
  112. });
  113. w.node_ids.push(ids::file_node_id(file_path));
  114. w.stack.push(Scope { row: 0, kind: "file", name: base_name.to_string() });
  115. w.visit_node(tree.root_node());
  116. w.flush_fn_ref_candidates();
  117. w.flush_value_refs(tree.root_node());
  118. w.stack.pop();
  119. let duration_ms = t0.elapsed().as_secs_f64() * 1000.0;
  120. let meta = build_meta(&w.tables, w.arena.len(), NONE_STR, duration_ms);
  121. Ok(EmitOut {
  122. meta,
  123. nodes: w.tables.nodes,
  124. edges: w.tables.edges,
  125. refs: w.tables.refs,
  126. arena: w.arena.into_vec(),
  127. })
  128. }
  129. impl<'t> Walker<'t> {
  130. fn text(&self, node: Node) -> &'t str {
  131. &self.src[node.byte_range()]
  132. }
  133. fn line_of(&self, node: Node) -> u32 {
  134. node.start_position().row as u32 + 1
  135. }
  136. fn col_of(&self, node: Node) -> u32 {
  137. util::col16(self.src, &self.line_starts, node.start_position().row, node.start_byte())
  138. }
  139. fn end_col_of(&self, node: Node) -> u32 {
  140. util::col16(self.src, &self.line_starts, node.end_position().row, node.end_byte())
  141. }
  142. fn top_row(&self) -> u32 {
  143. self.stack.last().map(|s| s.row).unwrap_or(0)
  144. }
  145. fn inside_class_like(&self) -> bool {
  146. self.stack
  147. .last()
  148. .map(|s| matches!(s.kind, "class" | "struct" | "interface" | "trait" | "enum" | "module"))
  149. .unwrap_or(false)
  150. }
  151. fn push_ref_at(&mut self, from_row: u32, name: &str, kind_code: u8, node: Node) {
  152. let name_ref = self.arena.put(name);
  153. self.tables.push_ref(&RefRow {
  154. from_idx: from_row,
  155. kind: kind_code,
  156. line: self.line_of(node),
  157. column: self.col_of(node),
  158. reference_name: name_ref,
  159. candidates: NONE_STR,
  160. from_id_str: NONE_STR,
  161. });
  162. if kind_code == edge_kind_index("imports").unwrap() {
  163. if util::simple_name().is_match(name) {
  164. self.imported_names.insert(name.to_string());
  165. } else if let Some(c) = util::qualified_import().captures(name) {
  166. self.imported_names.insert(c[1].to_string());
  167. }
  168. }
  169. }
  170. fn create_node(&mut self, kind: &'static str, name: &str, node: Node<'t>, extra: Extra) -> Option<u32> {
  171. if name.is_empty() {
  172. return None;
  173. }
  174. let start_line = self.line_of(node);
  175. let id = ids::node_id(self.file_path, kind, name, start_line);
  176. let end_line = node.end_position().row as u32 + 1;
  177. let qualified = {
  178. let mut parts: Vec<&str> = Vec::new();
  179. for s in &self.stack {
  180. if s.kind != "file" {
  181. parts.push(&s.name);
  182. }
  183. }
  184. let mut qn = parts.join("::");
  185. if !qn.is_empty() {
  186. qn.push_str("::");
  187. }
  188. qn.push_str(name);
  189. qn
  190. };
  191. let mut flags = BoolFlags::default();
  192. if let Some(v) = extra.is_async {
  193. flags.set(FLAG_IS_ASYNC, v);
  194. }
  195. if let Some(v) = extra.is_static {
  196. flags.set(FLAG_IS_STATIC, v);
  197. }
  198. let name_ref = self.arena.put(name);
  199. let qn_ref = self.arena.put(&qualified);
  200. let id_ref = self.arena.put(&id);
  201. let doc_ref = opt_str(&mut self.arena, extra.docstring.as_deref());
  202. let sig_ref = opt_str(&mut self.arena, extra.signature.as_deref());
  203. let row = self.tables.push_node(&NodeRow {
  204. kind: node_kind_index(kind).unwrap(),
  205. visibility: 0,
  206. flags,
  207. start_line,
  208. end_line,
  209. start_column: self.col_of(node),
  210. end_column: self.end_col_of(node),
  211. name: name_ref,
  212. qualified_name: qn_ref,
  213. id: id_ref,
  214. docstring: doc_ref,
  215. signature: sig_ref,
  216. decorators: NONE_STR,
  217. type_parameters: NONE_STR,
  218. return_type: NONE_STR,
  219. extra_json: NONE_STR,
  220. });
  221. self.node_ids.push(id);
  222. let parent_row = self.top_row();
  223. self.tables.push_edge(&EdgeRow {
  224. source_idx: parent_row,
  225. target_idx: row,
  226. kind: edge_kind_index("contains").unwrap(),
  227. provenance: 0,
  228. line: NONE,
  229. column: NONE,
  230. metadata_json: NONE_STR,
  231. source_id_str: NONE_STR,
  232. target_id_str: NONE_STR,
  233. });
  234. if kind == "function" || kind == "method" {
  235. self.defined_fn_names.insert(name.to_string());
  236. }
  237. // captureValueRefScope
  238. let target_kind_ok = kind == "constant" || kind == "variable";
  239. if target_kind_ok
  240. && util::utf16_len(name) >= 3
  241. && util::has_upper_or_underscore().is_match(name)
  242. {
  243. let parent_ok = self
  244. .stack
  245. .last()
  246. .map(|s| matches!(s.kind, "file" | "class" | "module" | "struct" | "enum"))
  247. .unwrap_or(false);
  248. if parent_ok {
  249. self.fs_values.insert(name.to_string(), row);
  250. *self.fs_value_counts.entry(name.to_string()).or_insert(0) += 1;
  251. }
  252. }
  253. if matches!(kind, "function" | "method" | "constant" | "variable") {
  254. self.value_scopes.push(ValueScope { row, node, name: name.to_string() });
  255. }
  256. Some(row)
  257. }
  258. fn extract_name(&self, node: Node) -> String {
  259. if let Some(name_node) = node.child_by_field_name("name") {
  260. return self.text(name_node).to_string();
  261. }
  262. for i in 0..node.named_child_count() {
  263. if let Some(c) = node.named_child(i) {
  264. if matches!(c.kind(), "identifier" | "type_identifier" | "simple_identifier" | "constant") {
  265. return self.text(c).to_string();
  266. }
  267. }
  268. }
  269. "<anonymous>".to_string()
  270. }
  271. /// pythonExtractor.getSignature: params + ` -> returnType`.
  272. fn signature_of(&self, node: Node) -> Option<String> {
  273. let params = node.child_by_field_name("parameters")?;
  274. let mut sig = self.text(params).to_string();
  275. if let Some(ret) = node.child_by_field_name("return_type") {
  276. sig.push_str(" -> ");
  277. sig.push_str(self.text(ret));
  278. }
  279. Some(sig)
  280. }
  281. /// pythonExtractor.isAsync: the PREVIOUS SIBLING token is `async`.
  282. fn is_async(&self, node: Node) -> bool {
  283. node.prev_sibling().map(|p| p.kind() == "async").unwrap_or(false)
  284. }
  285. /// pythonExtractor.isStatic: preceding decorator mentioning `staticmethod`.
  286. fn is_static(&self, node: Node) -> bool {
  287. if let Some(prev) = node.prev_named_sibling() {
  288. if prev.kind() == "decorator" {
  289. return self.text(prev).contains("staticmethod");
  290. }
  291. }
  292. false
  293. }
  294. // --- visitNode ------------------------------------------------------------
  295. fn visit_node(&mut self, node: Node<'t>) {
  296. let kind = node.kind();
  297. let mut skip_children = false;
  298. self.maybe_capture_fn_refs(node);
  299. if kind == "function_definition" {
  300. // functionTypes ∩ methodTypes: inside a class-like ⇒ method.
  301. if self.inside_class_like() {
  302. self.extract_method(node);
  303. } else {
  304. self.extract_function(node);
  305. }
  306. skip_children = true;
  307. } else if kind == "class_definition" {
  308. self.extract_class(node);
  309. skip_children = true;
  310. } else if kind == "assignment" && !self.inside_class_like() {
  311. self.extract_variable(node);
  312. self.scan_fn_ref_subtree(node, 0);
  313. skip_children = true;
  314. } else if kind == "import_statement" || kind == "import_from_statement" {
  315. self.extract_import(node);
  316. } else if kind == "call" {
  317. self.extract_call(node);
  318. }
  319. if !skip_children {
  320. for i in 0..node.named_child_count() {
  321. if let Some(c) = node.named_child(i) {
  322. self.visit_node(c);
  323. }
  324. }
  325. }
  326. }
  327. fn visit_function_body(&mut self, body: Node<'t>) {
  328. self.visit_for_calls_and_structure(body);
  329. }
  330. fn visit_for_calls_and_structure(&mut self, node: Node<'t>) {
  331. let kind = node.kind();
  332. self.maybe_capture_fn_refs(node);
  333. if kind == "call" {
  334. self.extract_call(node);
  335. }
  336. // Nested NAMED functions become their own nodes.
  337. if kind == "function_definition" {
  338. let name = self.extract_name(node);
  339. if name != "<anonymous>" {
  340. self.extract_function(node);
  341. return;
  342. }
  343. }
  344. if kind == "class_definition" {
  345. self.extract_class(node);
  346. return;
  347. }
  348. for i in 0..node.named_child_count() {
  349. if let Some(c) = node.named_child(i) {
  350. self.visit_for_calls_and_structure(c);
  351. }
  352. }
  353. }
  354. // --- extractors --------------------------------------------------------------
  355. fn extract_function(&mut self, node: Node<'t>) {
  356. let name = self.extract_name(node);
  357. if name == "<anonymous>" {
  358. if let Some(body) = node.child_by_field_name("body") {
  359. self.visit_function_body(body);
  360. }
  361. return;
  362. }
  363. let extra = Extra {
  364. docstring: preceding_docstring(node, self.src),
  365. signature: self.signature_of(node),
  366. is_async: Some(self.is_async(node)),
  367. is_static: Some(self.is_static(node)),
  368. };
  369. let Some(row) = self.create_node("function", &name, node, extra) else { return };
  370. // (python is not a TYPE_ANNOTATION language — no type refs)
  371. self.extract_decorators_for(node, row);
  372. self.stack.push(Scope { row, kind: "function", name });
  373. if let Some(body) = node.child_by_field_name("body") {
  374. self.visit_function_body(body);
  375. }
  376. self.stack.pop();
  377. }
  378. fn extract_method(&mut self, node: Node<'t>) {
  379. let name = self.extract_name(node);
  380. let extra = Extra {
  381. docstring: preceding_docstring(node, self.src),
  382. signature: self.signature_of(node),
  383. is_async: Some(self.is_async(node)),
  384. is_static: Some(self.is_static(node)),
  385. };
  386. let Some(row) = self.create_node("method", &name, node, extra) else { return };
  387. self.extract_decorators_for(node, row);
  388. self.stack.push(Scope { row, kind: "method", name });
  389. if let Some(body) = node.child_by_field_name("body") {
  390. self.visit_function_body(body);
  391. }
  392. self.stack.pop();
  393. }
  394. fn extract_class(&mut self, node: Node<'t>) {
  395. let name = self.extract_name(node);
  396. let extra = Extra {
  397. docstring: preceding_docstring(node, self.src),
  398. ..Extra::default()
  399. };
  400. let Some(row) = self.create_node("class", &name, node, extra) else { return };
  401. // Inheritance: `class Flask(Scaffold, Mixin):` — argument_list children.
  402. let extends_kind = edge_kind_index("extends").unwrap();
  403. for i in 0..node.named_child_count() {
  404. let Some(child) = node.named_child(i) else { continue };
  405. if child.kind() == "argument_list" {
  406. for j in 0..child.named_child_count() {
  407. let Some(arg) = child.named_child(j) else { continue };
  408. if matches!(arg.kind(), "identifier" | "attribute") {
  409. let name = self.text(arg).to_string();
  410. self.push_ref_at(row, &name, extends_kind, arg);
  411. }
  412. }
  413. }
  414. }
  415. self.extract_decorators_for(node, row);
  416. self.stack.push(Scope { row, kind: "class", name });
  417. let body = node.child_by_field_name("body").unwrap_or(node);
  418. for i in 0..body.named_child_count() {
  419. if let Some(c) = body.named_child(i) {
  420. self.visit_node(c);
  421. }
  422. }
  423. self.stack.pop();
  424. }
  425. /// extractVariable's python branch: `left = right` at module scope.
  426. fn extract_variable(&mut self, node: Node<'t>) {
  427. let docstring = preceding_docstring(node, self.src);
  428. let left = node.child_by_field_name("left").or_else(|| node.named_child(0));
  429. let right = node.child_by_field_name("right").or_else(|| node.named_child(1));
  430. let Some(left) = left else { return };
  431. if !matches!(left.kind(), "identifier" | "constant") {
  432. return;
  433. }
  434. let name = self.text(left).to_string();
  435. let signature = right.map(|r| util::init_signature(self.text(r)));
  436. // No isConst hook ⇒ always `variable` (UPPER_CASE constants included).
  437. self.create_node("variable", &name, node, Extra { docstring, signature, ..Extra::default() });
  438. }
  439. fn extract_import(&mut self, node: Node<'t>) {
  440. let import_text = self.text(node).trim().to_string();
  441. let imports_kind = edge_kind_index("imports").unwrap();
  442. if node.kind() == "import_from_statement" {
  443. // Hook path: module_name field → import node + module ref, then
  444. // per-name binding refs (emitPyFromImportRefs).
  445. let Some(module_node) = node.child_by_field_name("module_name") else { return };
  446. let module_name = self.text(module_node).to_string();
  447. if module_name.is_empty() {
  448. return;
  449. }
  450. self.create_node(
  451. "import",
  452. &module_name,
  453. node,
  454. Extra { signature: Some(import_text), ..Extra::default() },
  455. );
  456. let parent = self.top_row();
  457. self.push_ref_at(parent, &module_name.clone(), imports_kind, node);
  458. // emitPyFromImportRefs: one `imports` ref per imported name.
  459. for i in 0..node.named_child_count() {
  460. let Some(child) = node.named_child(i) else { continue };
  461. if child.start_byte() == module_node.start_byte()
  462. && child.end_byte() == module_node.end_byte()
  463. {
  464. continue;
  465. }
  466. if child.kind() == "wildcard_import" {
  467. continue;
  468. }
  469. let name_node = match child.kind() {
  470. "aliased_import" => child
  471. .child_by_field_name("alias")
  472. .or_else(|| child.child_by_field_name("name"))
  473. .or_else(|| child.named_child(0)),
  474. "dotted_name" => Some(child),
  475. _ => None,
  476. };
  477. let Some(name_node) = name_node else { continue };
  478. let raw = self.text(name_node);
  479. let local = raw.rsplit('.').next().unwrap_or("");
  480. if local.is_empty() {
  481. continue;
  482. }
  483. self.push_ref_at(parent, &local.to_string(), imports_kind, name_node);
  484. }
  485. return;
  486. }
  487. // import_statement: `import a.b, x as y` — one import node + module ref
  488. // per dotted name (the python multi-import branch).
  489. let parent = self.top_row();
  490. for i in 0..node.named_child_count() {
  491. let Some(child) = node.named_child(i) else { continue };
  492. if child.kind() == "dotted_name" {
  493. let name = self.text(child).to_string();
  494. self.create_node(
  495. "import",
  496. &name,
  497. node,
  498. Extra { signature: Some(import_text.clone()), ..Extra::default() },
  499. );
  500. self.push_ref_at(parent, &name, imports_kind, child);
  501. } else if child.kind() == "aliased_import" {
  502. let dotted = (0..child.named_child_count())
  503. .filter_map(|j| child.named_child(j))
  504. .find(|c| c.kind() == "dotted_name");
  505. if let Some(dotted) = dotted {
  506. let name = self.text(dotted).to_string();
  507. self.create_node(
  508. "import",
  509. &name,
  510. node,
  511. Extra { signature: Some(import_text.clone()), ..Extra::default() },
  512. );
  513. self.push_ref_at(parent, &name, imports_kind, dotted);
  514. }
  515. }
  516. }
  517. }
  518. /// extractCall — python `call` through the generic tail (attribute callees).
  519. fn extract_call(&mut self, node: Node<'t>) {
  520. if self.stack.is_empty() {
  521. return;
  522. }
  523. let func = node
  524. .child_by_field_name("function")
  525. .or_else(|| node.named_child(0));
  526. let mut callee_name = String::new();
  527. if let Some(func) = func {
  528. if func.kind() == "attribute" {
  529. // `property` and `field` fields don't exist on attribute —
  530. // the generic path falls back to namedChild(1) (the attr name).
  531. let property = func
  532. .child_by_field_name("property")
  533. .or_else(|| func.child_by_field_name("field"))
  534. .or_else(|| func.named_child(1));
  535. if let Some(property) = property {
  536. let method_name = self.text(property);
  537. let receiver = func
  538. .child_by_field_name("object")
  539. .or_else(|| func.child_by_field_name("operand"))
  540. .or_else(|| func.child_by_field_name("argument"))
  541. .or_else(|| func.named_child(0));
  542. if let Some(r) = receiver {
  543. if is_literal_receiver(r.kind()) {
  544. return;
  545. }
  546. }
  547. let recv_ident = receiver.filter(|r| {
  548. matches!(r.kind(), "identifier" | "simple_identifier" | "field_identifier")
  549. });
  550. if let Some(r) = recv_ident {
  551. let receiver_name = self.text(r);
  552. if !matches!(receiver_name, "self" | "this" | "cls" | "super") {
  553. callee_name = format!("{receiver_name}.{method_name}");
  554. } else {
  555. callee_name = method_name.to_string();
  556. }
  557. } else {
  558. callee_name = method_name.to_string();
  559. }
  560. }
  561. } else {
  562. callee_name = self.text(func).to_string();
  563. }
  564. }
  565. if !callee_name.is_empty() {
  566. if let Some(c) = util::paren_conversion().captures(&callee_name) {
  567. callee_name = c[1].to_string();
  568. }
  569. let from = self.top_row();
  570. self.push_ref_at(from, &callee_name.clone(), edge_kind_index("calls").unwrap(), node);
  571. }
  572. }
  573. /// extractDecoratorsFor — python decorators are PRECEDING SIBLINGS inside
  574. /// decorated_definition. Only bare-identifier decorators yield a target
  575. /// (python's `call` kind isn't `call_expression`, and `attribute` isn't in
  576. /// the target-kind list — mirrored exactly).
  577. fn extract_decorators_for(&mut self, decl: Node<'t>, decorated_row: u32) {
  578. for i in 0..decl.named_child_count() {
  579. if let Some(child) = decl.named_child(i) {
  580. self.consider_decorator(child, decorated_row);
  581. }
  582. }
  583. let Some(parent) = decl.parent() else { return };
  584. let decl_start = decl.start_byte();
  585. let mut decl_idx: isize = -1;
  586. for i in 0..parent.named_child_count() {
  587. if let Some(sib) = parent.named_child(i) {
  588. if sib.start_byte() == decl_start {
  589. decl_idx = i as isize;
  590. break;
  591. }
  592. }
  593. }
  594. if decl_idx > 0 {
  595. let mut j = decl_idx - 1;
  596. while j >= 0 {
  597. let Some(sib) = parent.named_child(j as usize) else {
  598. j -= 1;
  599. continue;
  600. };
  601. if !matches!(sib.kind(), "decorator" | "annotation" | "marker_annotation") {
  602. break;
  603. }
  604. self.consider_decorator(sib, decorated_row);
  605. j -= 1;
  606. }
  607. }
  608. }
  609. fn consider_decorator(&mut self, n: Node<'t>, decorated_row: u32) {
  610. if !matches!(n.kind(), "decorator" | "annotation" | "marker_annotation" | "attribute") {
  611. return;
  612. }
  613. let mut target: Option<Node> = None;
  614. for i in 0..n.named_child_count() {
  615. let Some(child) = n.named_child(i) else { continue };
  616. if child.kind() == "call_expression" {
  617. target = child.child_by_field_name("function").or_else(|| child.named_child(0));
  618. if target.is_some() {
  619. break;
  620. }
  621. }
  622. if matches!(
  623. child.kind(),
  624. "identifier" | "member_expression" | "scoped_identifier" | "navigation_expression"
  625. | "user_type" | "type_identifier"
  626. ) {
  627. target = Some(child);
  628. break;
  629. }
  630. }
  631. let Some(target) = target else { return };
  632. let mut name = self.text(target).to_string();
  633. if let Some(lt) = name.find('<') {
  634. if lt > 0 {
  635. name.truncate(lt);
  636. }
  637. }
  638. let last_dot = name
  639. .rfind('.')
  640. .map(|i| i as isize)
  641. .unwrap_or(-1)
  642. .max(name.rfind("::").map(|i| i as isize).unwrap_or(-1));
  643. if last_dot >= 0 {
  644. name = name[(last_dot as usize + 1)..].to_string();
  645. if name.starts_with(':') || name.starts_with('.') {
  646. name.remove(0);
  647. }
  648. }
  649. let name = name.trim().to_string();
  650. if name.is_empty() {
  651. return;
  652. }
  653. self.push_ref_at(decorated_row, &name, edge_kind_index("decorates").unwrap(), n);
  654. }
  655. // --- fn refs (PYTHON_SPEC) ------------------------------------------------------
  656. fn maybe_capture_fn_refs(&mut self, node: Node<'t>) {
  657. let (mode, field): (&str, &str) = match node.kind() {
  658. "argument_list" => ("args", ""),
  659. "assignment" => ("rhs", "right"),
  660. "keyword_argument" => ("value", "value"),
  661. "pair" => ("value", "value"),
  662. "list" => ("list", ""),
  663. _ => return,
  664. };
  665. if self.stack.is_empty() {
  666. return;
  667. }
  668. let from = self.top_row();
  669. let mut values: Vec<Node> = Vec::new();
  670. match mode {
  671. "args" | "list" => {
  672. for i in 0..node.named_child_count() {
  673. if let Some(c) = node.named_child(i) {
  674. values.push(c);
  675. }
  676. }
  677. }
  678. "rhs" => {
  679. if let Some(rhs) = node.child_by_field_name(field) {
  680. let lhs_text = node
  681. .child_by_field_name("left")
  682. .map(|l| self.text(l))
  683. .unwrap_or("");
  684. let lhs_last = util::lhs_last_name()
  685. .captures(lhs_text)
  686. .and_then(|c| c.get(1))
  687. .map(|m| m.as_str());
  688. if !(lhs_last.is_some() && lhs_last == Some(self.text(rhs).trim())) {
  689. values.push(rhs);
  690. }
  691. }
  692. }
  693. _ => {
  694. if let Some(v) = node.child_by_field_name(field) {
  695. values.push(v);
  696. }
  697. }
  698. }
  699. for v in values {
  700. let (name, anchor) = match v.kind() {
  701. "identifier" => (self.text(v).to_string(), v),
  702. // `self.handle_click` — object EXACTLY `self`; BARE attr name.
  703. "attribute" => {
  704. let obj = v.child_by_field_name("object");
  705. let attr = v.child_by_field_name("attribute");
  706. match (obj, attr) {
  707. (Some(o), Some(a))
  708. if o.kind() == "identifier" && self.text(o) == "self" =>
  709. {
  710. (self.text(a).to_string(), a)
  711. }
  712. _ => continue,
  713. }
  714. }
  715. _ => continue,
  716. };
  717. if name.is_empty() || is_stoplisted(&name) {
  718. continue;
  719. }
  720. let p = anchor.start_position();
  721. self.fn_ref_cands.push(Cand {
  722. from,
  723. name,
  724. line: p.row as u32 + 1,
  725. column_byte: anchor.start_byte(),
  726. row: p.row,
  727. });
  728. }
  729. }
  730. fn scan_fn_ref_subtree(&mut self, node: Node<'t>, depth: u32) {
  731. if depth > 12 {
  732. return;
  733. }
  734. // Halts at functionTypes ∪ the fixed arrow/lambda list — python's
  735. // `lambda` kind is NOT in that list (mirrored).
  736. if depth > 0
  737. && matches!(
  738. node.kind(),
  739. "function_definition" | "arrow_function" | "function_expression" | "lambda_literal"
  740. | "lambda_expression"
  741. )
  742. {
  743. return;
  744. }
  745. self.maybe_capture_fn_refs(node);
  746. for i in 0..node.named_child_count() {
  747. if let Some(c) = node.named_child(i) {
  748. self.scan_fn_ref_subtree(c, depth + 1);
  749. }
  750. }
  751. }
  752. fn flush_fn_ref_candidates(&mut self) {
  753. let cands = std::mem::take(&mut self.fn_ref_cands);
  754. if cands.is_empty() || util::is_generated_file(self.file_path) {
  755. return;
  756. }
  757. let mut seen: HashSet<(String, String)> = HashSet::new();
  758. for c in cands {
  759. if !c.name.starts_with("this.")
  760. && !c.name.contains("::")
  761. && !self.defined_fn_names.contains(&c.name)
  762. && !self.imported_names.contains(&c.name)
  763. {
  764. continue;
  765. }
  766. if !seen.insert((self.node_ids[c.from as usize].clone(), c.name.clone())) {
  767. continue;
  768. }
  769. let column = util::col16(self.src, &self.line_starts, c.row, c.column_byte);
  770. let name_ref = self.arena.put(&c.name);
  771. self.tables.push_ref(&RefRow {
  772. from_idx: c.from,
  773. kind: FUNCTION_REF_CODE,
  774. line: c.line,
  775. column,
  776. reference_name: name_ref,
  777. candidates: NONE_STR,
  778. from_id_str: NONE_STR,
  779. });
  780. }
  781. }
  782. // --- value refs -------------------------------------------------------------------
  783. fn flush_value_refs(&mut self, root: Node<'t>) {
  784. let scopes = std::mem::take(&mut self.value_scopes);
  785. let mut targets = std::mem::take(&mut self.fs_values);
  786. let counts = std::mem::take(&mut self.fs_value_counts);
  787. if std::env::var("CODEGRAPH_VALUE_REFS").as_deref() == Ok("0") {
  788. return;
  789. }
  790. if targets.is_empty() || scopes.is_empty() || util::is_generated_file(self.file_path) {
  791. return;
  792. }
  793. // Shadow prune — python's declarator shape is `assignment`.
  794. let mut decl_counts: HashMap<&str, u32> = HashMap::new();
  795. let mut dstack: Vec<Node> = vec![root];
  796. let mut dvisited = 0usize;
  797. while let Some(n) = dstack.pop() {
  798. if dvisited >= MAX_VALUE_REF_NODES {
  799. break;
  800. }
  801. dvisited += 1;
  802. if n.kind() == "assignment" {
  803. let left = n
  804. .child_by_field_name("left")
  805. .or_else(|| n.child_by_field_name("pattern"))
  806. .or_else(|| n.named_child(0));
  807. if let Some(left) = left {
  808. if left.kind() == "identifier" {
  809. let nm = self.text(left);
  810. if targets.contains_key(nm) {
  811. *decl_counts.entry(nm).or_insert(0) += 1;
  812. }
  813. } else {
  814. for i in 0..left.named_child_count() {
  815. if let Some(c) = left.named_child(i) {
  816. if c.kind() == "identifier" {
  817. let nm = self.text(c);
  818. if targets.contains_key(nm) {
  819. *decl_counts.entry(nm).or_insert(0) += 1;
  820. }
  821. }
  822. }
  823. }
  824. }
  825. }
  826. }
  827. for i in 0..n.named_child_count() {
  828. if let Some(c) = n.named_child(i) {
  829. dstack.push(c);
  830. }
  831. }
  832. }
  833. let shadowed: Vec<String> = decl_counts
  834. .iter()
  835. .filter(|(nm, c)| **c > counts.get(**nm).copied().unwrap_or(1))
  836. .map(|(nm, _)| nm.to_string())
  837. .collect();
  838. for nm in shadowed {
  839. targets.remove(&nm);
  840. }
  841. if targets.is_empty() {
  842. return;
  843. }
  844. let refs_kind = edge_kind_index("references").unwrap();
  845. for scope in &scopes {
  846. let mut seen: HashSet<&str> = HashSet::new();
  847. let mut stack: Vec<Node> = vec![scope.node];
  848. let mut visited = 0usize;
  849. while let Some(n) = stack.pop() {
  850. if visited >= MAX_VALUE_REF_NODES {
  851. break;
  852. }
  853. visited += 1;
  854. if matches!(n.kind(), "identifier" | "constant" | "name" | "simple_identifier") {
  855. let ref_name = self.text(n);
  856. if let Some(&target_row) = targets.get(ref_name) {
  857. let target_id = self.node_ids[target_row as usize].as_str();
  858. if target_id != self.node_ids[scope.row as usize]
  859. && ref_name != scope.name
  860. && !seen.contains(&target_id)
  861. {
  862. seen.insert(target_id);
  863. let meta = self.arena.put(r#"{"valueRef":true}"#);
  864. self.tables.push_edge(&EdgeRow {
  865. source_idx: scope.row,
  866. target_idx: target_row,
  867. kind: refs_kind,
  868. provenance: 0,
  869. line: NONE,
  870. column: NONE,
  871. metadata_json: meta,
  872. source_id_str: NONE_STR,
  873. target_id_str: NONE_STR,
  874. });
  875. }
  876. }
  877. }
  878. for i in 0..n.named_child_count() {
  879. if let Some(c) = n.named_child(i) {
  880. stack.push(c);
  881. }
  882. }
  883. }
  884. }
  885. }
  886. }
  887. /// NAME_STOPLIST (function-ref.ts).
  888. fn is_stoplisted(name: &str) -> bool {
  889. matches!(
  890. name,
  891. "this" | "self" | "super" | "null" | "nil" | "true" | "false" | "undefined" | "new"
  892. | "NULL" | "nullptr" | "None"
  893. )
  894. }
  895. /// LITERAL_RECEIVER_TYPES membership (shared table; python names among them).
  896. fn is_literal_receiver(kind: &str) -> bool {
  897. matches!(
  898. kind,
  899. "string" | "string_literal" | "interpreted_string_literal" | "raw_string_literal"
  900. | "template_string" | "concatenated_string" | "formatted_string" | "f_string"
  901. | "line_string_literal" | "string_content" | "heredoc_body"
  902. | "number" | "number_literal" | "integer" | "integer_literal" | "float"
  903. | "float_literal" | "int_literal" | "decimal_integer_literal" | "real_literal"
  904. | "char_literal" | "character_literal" | "rune_literal" | "regex" | "regex_literal"
  905. | "true" | "false" | "boolean_literal" | "bool_literal" | "none" | "null" | "nil"
  906. | "null_literal" | "undefined"
  907. | "list" | "list_literal" | "array" | "array_literal" | "array_creation_expression"
  908. | "dictionary" | "dict_literal" | "object" | "tuple" | "set"
  909. )
  910. }
  911. fn opt_str(arena: &mut Arena, s: Option<&str>) -> StrRef {
  912. match s {
  913. Some(s) => arena.put(s),
  914. None => NONE_STR,
  915. }
  916. }