1use crate::imports::python_import_infos_from_node;
2use crate::syntax::{PythonOverloadDecoratorBindings, expression_name_node};
3use brokk_bifrost_core::analyzer::fq_name::{FqName, SegmentId, SegmentKind, segment_interner};
4use brokk_bifrost_core::analyzer::model::{
5 CodeUnitType, DispatchExtensibility, ParameterMetadata, SignatureMetadata,
6};
7use brokk_bifrost_core::analyzer::parsed_file::ParsedFile;
8use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
9use brokk_bifrost_core::analyzer::{CodeUnit, ProjectFile};
10use brokk_bifrost_core::hash::HashSet;
11use brokk_bifrost_core::path_normalization::NormalizePath;
12use brokk_bifrost_core::text_utils::{compute_line_starts, find_line_index_for_offset};
13use std::path::{Path, PathBuf};
14use tree_sitter::{Node, Parser, Tree};
15
16fn py_segment(text: &str, kind: SegmentKind) -> SegmentId {
18 segment_interner().intern(text, kind)
19}
20
21pub fn python_module_fq(file: &ProjectFile) -> FqName {
34 python_module_fq_from_components(&python_module_components(file))
35}
36
37fn python_module_fq_from_components(components: &[String]) -> FqName {
38 let mut fq = FqName::new();
39 for component in components {
40 fq.push(py_segment(component, SegmentKind::Package));
41 }
42 fq
43}
44
45fn python_module_components(file: &ProjectFile) -> Vec<String> {
46 let mut components = python_package_components_for_file(file);
47 let module_name = file
48 .rel_path()
49 .file_stem()
50 .and_then(|stem| stem.to_str())
51 .unwrap_or_default();
52 if module_name != "__init__" || components.is_empty() {
53 components.push(module_name.to_string());
54 }
55 components
56}
57
58fn python_package_components_for_file(file: &ProjectFile) -> Vec<String> {
59 let Some(parent_rel) = file.rel_path().parent() else {
60 return Vec::new();
61 };
62 if parent_rel.as_os_str().is_empty() {
63 return Vec::new();
64 }
65
66 if let Some(import_root_rel) = python_configured_import_root(file, parent_rel)
67 && let Ok(relative_package) = parent_rel.strip_prefix(import_root_rel)
68 {
69 return path_components(relative_package);
70 }
71
72 let mut effective_package_root_rel: Option<&Path> = None;
73 let mut current_rel = Some(parent_rel);
74 while let Some(path) = current_rel {
75 if file.root().join(path).join("__init__.py").exists() {
76 effective_package_root_rel = Some(path);
77 }
78 current_rel = path.parent();
79 }
80
81 let relative_package = match effective_package_root_rel {
82 Some(package_root_rel) => package_root_rel
83 .parent()
84 .and_then(|import_root_rel| parent_rel.strip_prefix(import_root_rel).ok())
85 .unwrap_or(parent_rel),
86 None => parent_rel,
87 };
88 path_components(relative_package)
89}
90
91fn python_configured_import_root(file: &ProjectFile, parent_rel: &Path) -> Option<PathBuf> {
96 let mut manifest_dir_rel = Some(parent_rel);
97 while let Some(directory) = manifest_dir_rel {
98 let manifest_dir = file.root().join(directory);
99 let mut roots = setuptools_where_entries(&manifest_dir.join("pyproject.toml"))
100 .iter()
101 .map(|entry| manifest_dir.join(entry).normalize())
102 .filter_map(|root| root.strip_prefix(file.root()).ok().map(Path::to_path_buf))
103 .filter(|root| parent_rel.starts_with(root))
104 .collect::<Vec<_>>();
105 roots.sort_by_key(|root| root.components().count());
106 if let Some(root) = roots.pop() {
107 return Some(root);
108 }
109 manifest_dir_rel = directory.parent();
110 }
111 None
112}
113
114struct ManifestWhereEntries {
119 len: u64,
120 modified: Option<std::time::SystemTime>,
121 entries: Vec<String>,
122}
123
124fn setuptools_where_entries(manifest: &Path) -> Vec<String> {
135 static MEMO: std::sync::OnceLock<
136 std::sync::RwLock<std::collections::HashMap<PathBuf, ManifestWhereEntries>>,
137 > = std::sync::OnceLock::new();
138 let memo = MEMO.get_or_init(Default::default);
139
140 let Ok(metadata) = std::fs::metadata(manifest) else {
141 return Vec::new();
142 };
143 let (len, modified) = (metadata.len(), metadata.modified().ok());
144 if let Some(cached) = memo.read().expect("manifest memo").get(manifest)
145 && cached.len == len
146 && cached.modified == modified
147 {
148 return cached.entries.clone();
149 }
150
151 let entries = parse_setuptools_where_entries(manifest);
152 memo.write().expect("manifest memo").insert(
153 manifest.to_path_buf(),
154 ManifestWhereEntries {
155 len,
156 modified,
157 entries: entries.clone(),
158 },
159 );
160 entries
161}
162
163fn parse_setuptools_where_entries(manifest: &Path) -> Vec<String> {
164 let Ok(source) = std::fs::read_to_string(manifest) else {
165 return Vec::new();
166 };
167 let Ok(document) = source.parse::<toml::Value>() else {
168 return Vec::new();
169 };
170 document
171 .get("tool")
172 .and_then(|tool| tool.get("setuptools"))
173 .and_then(|setuptools| setuptools.get("packages"))
174 .and_then(|packages| packages.get("find"))
175 .and_then(|find| find.get("where"))
176 .and_then(toml::Value::as_array)
177 .map(|entries| {
178 entries
179 .iter()
180 .filter_map(toml::Value::as_str)
181 .map(str::to_string)
182 .collect()
183 })
184 .unwrap_or_default()
185}
186
187fn path_components(path: &Path) -> Vec<String> {
188 path.components()
189 .map(|component| component.as_os_str().to_string_lossy().to_string())
190 .filter(|component| !component.is_empty())
191 .collect()
192}
193
194pub fn python_is_decorated_function_boundary(node: Node<'_>) -> bool {
195 if node.kind() != "decorated_definition" {
196 return false;
197 }
198 let mut cursor = node.walk();
199 node.named_children(&mut cursor)
200 .any(|child| child.kind() == "function_definition")
201}
202
203#[derive(Clone)]
204pub struct Scope {
205 kind: ScopeKind,
206 path: String,
207 fq: FqName,
214 code_unit: Option<CodeUnit>,
215 method_receiver: Option<String>,
216}
217
218#[derive(Clone, Copy, PartialEq, Eq)]
219enum ScopeKind {
220 Class,
221 Function,
222}
223
224pub struct PythonVisitor<'a> {
225 pub file: &'a ProjectFile,
226 pub source: &'a str,
227 pub package_name: &'a str,
228 module_fq: &'a FqName,
229 pub parsed: &'a mut ParsedFile,
230 pub module: Option<CodeUnit>,
231 pub overload_decorators: &'a PythonOverloadDecoratorBindings,
232}
233
234struct PythonContainer<'tree> {
235 node: Node<'tree>,
236 scope: Vec<Scope>,
237 module_control_depth: usize,
238}
239
240enum PythonWork<'tree> {
241 Container(PythonContainer<'tree>),
242 Statement {
243 node: Node<'tree>,
244 scope: Vec<Scope>,
245 module_control_depth: usize,
246 },
247}
248
249impl<'a> PythonVisitor<'a> {
250 pub fn visit_container(
251 &mut self,
252 node: Node<'_>,
253 scope: &[Scope],
254 module_control_depth: usize,
255 ) {
256 let mut stack = vec![PythonWork::Container(PythonContainer {
257 node,
258 scope: scope.to_vec(),
259 module_control_depth,
260 })];
261 while let Some(work) = stack.pop() {
262 match work {
263 PythonWork::Container(container) => {
264 let mut cursor = container.node.walk();
265 let children = container
266 .node
267 .named_children(&mut cursor)
268 .collect::<Vec<_>>();
269 for child in children.into_iter().rev() {
270 stack.push(PythonWork::Statement {
271 node: child,
272 scope: container.scope.clone(),
273 module_control_depth: container.module_control_depth,
274 });
275 }
276 }
277 PythonWork::Statement {
278 node,
279 scope,
280 module_control_depth,
281 } => self.visit_statement(node, &scope, module_control_depth, &mut stack),
282 }
283 }
284 }
285
286 fn visit_statement<'tree>(
287 &mut self,
288 node: Node<'tree>,
289 scope: &[Scope],
290 module_control_depth: usize,
291 stack: &mut Vec<PythonWork<'tree>>,
292 ) {
293 match node.kind() {
294 "decorated_definition" => {
295 if let Some(definition) = node.child_by_field_name("definition") {
296 self.visit_definition(
297 definition,
298 Some(node),
299 scope,
300 module_control_depth,
301 stack,
302 );
303 }
304 }
305 "class_definition" | "function_definition" => {
306 self.visit_definition(node, None, scope, module_control_depth, stack)
307 }
308 "expression_statement" => {
309 self.visit_expression_statement(node, scope, module_control_depth)
310 }
311 "import_statement" | "import_from_statement" => self.visit_import_statement(node),
312 "if_statement" | "try_statement" | "with_statement" | "for_statement"
313 | "while_statement" => {
314 let next_depth = if scope.is_empty() {
315 module_control_depth + 1
316 } else {
317 module_control_depth
318 };
319 stack.push(PythonWork::Container(PythonContainer {
320 node,
321 scope: scope.to_vec(),
322 module_control_depth: next_depth,
323 }));
324 }
325 "elif_clause" | "else_clause" | "except_clause" | "finally_clause" => {
326 stack.push(PythonWork::Container(PythonContainer {
327 node,
328 scope: scope.to_vec(),
329 module_control_depth,
330 }));
331 }
332 "block" | "module" => stack.push(PythonWork::Container(PythonContainer {
333 node,
334 scope: scope.to_vec(),
335 module_control_depth,
336 })),
337 _ => {}
338 }
339 }
340
341 fn visit_definition<'tree>(
342 &mut self,
343 definition: Node<'tree>,
344 wrapper: Option<Node<'tree>>,
345 scope: &[Scope],
346 module_control_depth: usize,
347 stack: &mut Vec<PythonWork<'tree>>,
348 ) {
349 match definition.kind() {
350 "class_definition" => self.visit_class_definition(
351 definition,
352 wrapper.unwrap_or(definition),
353 scope,
354 module_control_depth,
355 stack,
356 ),
357 "function_definition" => self.visit_function_definition(
358 definition,
359 wrapper.unwrap_or(definition),
360 scope,
361 module_control_depth,
362 stack,
363 ),
364 _ => {}
365 }
366 }
367
368 fn visit_class_definition<'tree>(
369 &mut self,
370 node: Node<'tree>,
371 range_node: Node<'tree>,
372 scope: &[Scope],
373 module_control_depth: usize,
374 stack: &mut Vec<PythonWork<'tree>>,
375 ) {
376 let Some(name_node) = node.child_by_field_name("name") else {
377 return;
378 };
379 let name = py_node_text(name_node, self.source).trim();
380 if name.is_empty() {
381 return;
382 }
383
384 let capture = !scope.is_empty() || module_control_depth <= 1;
385
386 let short_name = scope
387 .last()
388 .map(|parent| format!("{}${name}", parent.path))
389 .unwrap_or_else(|| name.to_string());
390 let fq = match scope.last() {
396 Some(parent) => parent
397 .fq
398 .clone()
399 .with_pushed(py_segment(name, SegmentKind::Nested)),
400 None => self
401 .module_fq
402 .clone()
403 .with_pushed(py_segment(name, SegmentKind::Type)),
404 };
405 let code_unit = CodeUnit::new_fq(
406 self.file.clone(),
407 CodeUnitType::Class,
408 self.package_name.to_string(),
409 short_name.clone(),
410 fq.clone(),
411 );
412 if capture {
413 self.parsed
414 .replace_code_unit(code_unit.clone(), range_node, self.source, None, None);
415 self.parsed.add_signature(
416 code_unit.clone(),
417 python_class_signature(range_node, self.source),
418 );
419 if let Some(module) = &self.module
420 && scope.is_empty()
421 {
422 self.parsed.add_child(module.clone(), code_unit.clone());
423 }
424 if let Some(parent) = scope.last()
425 && let Some(parent_cu) = &parent.code_unit
426 {
427 self.parsed.add_child(parent_cu.clone(), code_unit.clone());
428 }
429 self.parsed.set_raw_supertypes(
430 code_unit.clone(),
431 extract_python_supertypes(node, self.source),
432 );
433 }
434
435 let mut next_scope = scope.to_vec();
436 if capture {
437 next_scope.push(Scope {
438 kind: ScopeKind::Class,
439 path: short_name,
440 fq,
441 code_unit: Some(code_unit),
442 method_receiver: None,
443 });
444 }
445 if let Some(body) = node.child_by_field_name("body") {
446 stack.push(PythonWork::Container(PythonContainer {
447 node: body,
448 scope: next_scope,
449 module_control_depth,
450 }));
451 }
452 }
453
454 fn visit_function_definition<'tree>(
455 &mut self,
456 node: Node<'tree>,
457 range_node: Node<'tree>,
458 scope: &[Scope],
459 module_control_depth: usize,
460 stack: &mut Vec<PythonWork<'tree>>,
461 ) {
462 let Some(name_node) = node.child_by_field_name("name") else {
463 return;
464 };
465 let name = py_node_text(name_node, self.source).trim();
466 if name.is_empty() {
467 return;
468 }
469
470 let capture = !python_is_property_mutator(range_node, self.source)
471 && ((scope.is_empty() && module_control_depth <= 1)
472 || scope
473 .last()
474 .is_some_and(|parent| parent.kind == ScopeKind::Class));
475 let short_name = if let Some(parent) = scope.last() {
476 match parent.kind {
477 ScopeKind::Class => format!("{}.{}", parent.path, name),
478 ScopeKind::Function => format!("{}${name}", parent.path),
479 }
480 } else {
481 name.to_string()
482 };
483 let fq = if let Some(parent) = scope.last() {
488 match parent.kind {
489 ScopeKind::Class => parent
490 .fq
491 .clone()
492 .with_pushed(py_segment(name, SegmentKind::Member)),
493 ScopeKind::Function => parent
494 .fq
495 .clone()
496 .with_pushed(py_segment(name, SegmentKind::Nested)),
497 }
498 } else {
499 self.module_fq
500 .clone()
501 .with_pushed(py_segment(name, SegmentKind::Member))
502 };
503
504 if capture {
505 let code_unit_type = if python_function_has_decorator(node, self.source, "property") {
506 CodeUnitType::Field
507 } else {
508 CodeUnitType::Function
509 };
510 let signature = node
511 .child_by_field_name("parameters")
512 .map(|parameters| py_node_text(parameters, self.source).trim().to_string());
513 let code_unit = CodeUnit::with_signature_and_fq(
514 self.file.clone(),
515 code_unit_type,
516 self.package_name.to_string(),
517 short_name.clone(),
518 signature,
519 false,
520 fq.clone(),
521 );
522 self.parsed
523 .replace_code_unit(code_unit.clone(), range_node, self.source, None, None);
524 let signature = python_function_signature(range_node, self.source);
525 self.parsed.add_signature_with_metadata(
526 code_unit.clone(),
527 python_signature_metadata(signature, node, self.source).with_declaration_only(
528 self.overload_decorators
529 .decorates_as_overload(node, self.source),
530 ),
531 );
532 if let Some(module) = &self.module
533 && scope.is_empty()
534 {
535 self.parsed.add_child(module.clone(), code_unit.clone());
536 }
537 if let Some(parent) = scope.last()
538 && parent.kind == ScopeKind::Class
539 && let Some(parent_cu) = &parent.code_unit
540 {
541 self.parsed.add_child(parent_cu.clone(), code_unit.clone());
542 }
543 let scope_code_unit = Some(code_unit);
544 let mut next_scope = scope.to_vec();
545 next_scope.push(Scope {
546 kind: ScopeKind::Function,
547 path: short_name,
548 fq,
549 code_unit: scope_code_unit,
550 method_receiver: scope
551 .last()
552 .is_some_and(|parent| parent.kind == ScopeKind::Class)
553 .then(|| python_instance_method_receiver_name(node, self.source))
554 .flatten(),
555 });
556 if let Some(body) = node.child_by_field_name("body") {
557 stack.push(PythonWork::Container(PythonContainer {
558 node: body,
559 scope: next_scope,
560 module_control_depth,
561 }));
562 }
563 return;
564 }
565
566 let mut next_scope = scope.to_vec();
567 next_scope.push(Scope {
568 kind: ScopeKind::Function,
569 path: short_name,
570 fq,
571 code_unit: None,
572 method_receiver: None,
573 });
574 if let Some(body) = node.child_by_field_name("body") {
575 stack.push(PythonWork::Container(PythonContainer {
576 node: body,
577 scope: next_scope,
578 module_control_depth,
579 }));
580 }
581 }
582
583 fn visit_expression_statement(
584 &mut self,
585 node: Node<'_>,
586 scope: &[Scope],
587 module_control_depth: usize,
588 ) {
589 let Some(assignment) = node.named_child(0) else {
590 return;
591 };
592 if assignment.kind() != "assignment" {
593 return;
594 }
595 let Some(left) = assignment.child_by_field_name("left") else {
596 return;
597 };
598 self.visit_instance_attribute_assignment(left, scope);
599 let names = collect_assigned_names(left, self.source);
600 for name in names {
601 let (short_name, fq) = if let Some(parent) = scope.last() {
602 if parent.kind != ScopeKind::Class {
603 continue;
604 }
605 (
606 format!("{}.{}", parent.path, name),
607 parent
608 .fq
609 .clone()
610 .with_pushed(py_segment(&name, SegmentKind::Member)),
611 )
612 } else if module_control_depth <= 1 {
613 (
614 name.clone(),
615 self.module_fq
616 .clone()
617 .with_pushed(py_segment(&name, SegmentKind::Member)),
618 )
619 } else {
620 continue;
621 };
622 let code_unit = CodeUnit::new_fq(
623 self.file.clone(),
624 CodeUnitType::Field,
625 self.package_name.to_string(),
626 short_name,
627 fq,
628 );
629 if scope
630 .last()
631 .is_some_and(|parent| parent.kind == ScopeKind::Class)
632 {
633 self.parsed
637 .add_code_unit(code_unit.clone(), node, self.source, None, None);
638 } else {
639 self.parsed
640 .replace_code_unit(code_unit.clone(), node, self.source, None, None);
641 }
642 self.parsed.add_signature(
643 code_unit.clone(),
644 py_node_text(node, self.source).trim().to_string(),
645 );
646 if let Some(module) = &self.module
647 && scope.is_empty()
648 {
649 self.parsed.add_child(module.clone(), code_unit.clone());
650 }
651 if let Some(parent) = scope.last()
652 && parent.kind == ScopeKind::Class
653 && let Some(parent_cu) = &parent.code_unit
654 {
655 self.parsed.add_child(parent_cu.clone(), code_unit);
656 }
657 }
658 }
659
660 fn visit_instance_attribute_assignment(&mut self, left: Node<'_>, scope: &[Scope]) {
661 let Some(function) = scope
662 .last()
663 .filter(|scope| scope.kind == ScopeKind::Function)
664 else {
665 return;
666 };
667 let Some(receiver) = function.method_receiver.as_deref() else {
668 return;
669 };
670 let Some(parent) = scope
671 .get(scope.len().saturating_sub(2))
672 .filter(|scope| scope.kind == ScopeKind::Class)
673 else {
674 return;
675 };
676 let Some(parent_cu) = parent.code_unit.clone() else {
677 return;
678 };
679 for (name, node) in collect_self_assigned_attributes(left, self.source, receiver) {
680 let code_unit = CodeUnit::new_fq(
681 self.file.clone(),
682 CodeUnitType::Field,
683 self.package_name.to_string(),
684 format!("{}.{}", parent.path, name),
685 parent
686 .fq
687 .clone()
688 .with_pushed(py_segment(&name, SegmentKind::Member)),
689 );
690 if !self.parsed.contains_declaration(&code_unit) {
691 self.parsed.replace_code_unit(
692 code_unit.clone(),
693 node,
694 self.source,
695 Some(parent_cu.clone()),
696 Some(parent_cu.clone()),
697 );
698 }
699 self.parsed.add_signature(
700 code_unit.clone(),
701 py_node_text(left, self.source).trim().to_string(),
702 );
703 }
704 }
705
706 fn visit_import_statement(&mut self, node: Node<'_>) {
707 for info in python_import_infos_from_node(node, self.source) {
708 self.parsed.imports.push(info);
709 }
710 }
711}
712
713pub fn parse_python_file(file: &ProjectFile, source: &str, tree: &Tree) -> ParsedFile {
717 let module_components = python_module_components(file);
718 let module_name = module_components.join(".");
719 let module_fq = python_module_fq_from_components(&module_components);
720 let mut parsed = ParsedFile::new(module_name.clone());
721 let root = tree.root_node();
722
723 collect_python_identifiers(root, source, &mut parsed.type_identifiers);
724
725 let module_code_unit = module_code_unit_from_fq(file, &module_components, module_fq.clone());
726 if let Some(module) = module_code_unit.clone() {
727 parsed.add_code_unit(module, root, source, None, None);
728 }
729
730 let overload_decorators = PythonOverloadDecoratorBindings::collect(root, source);
731 let mut visitor = PythonVisitor {
732 file,
733 source,
734 package_name: &module_name,
735 module_fq: &module_fq,
736 parsed: &mut parsed,
737 module: module_code_unit,
738 overload_decorators: &overload_decorators,
739 };
740 visitor.visit_container(root, &[], 0);
741
742 parsed
743}
744
745pub fn py_node_text<'a>(node: Node<'_>, source: &'a str) -> &'a str {
746 brokk_bifrost_core::analyzer::common::node_source_text(node, source)
747}
748
749pub fn python_module_name(file: &ProjectFile) -> String {
750 python_module_components(file).join(".")
751}
752
753pub fn module_code_unit(file: &ProjectFile, module_fq: &str) -> Option<CodeUnit> {
754 if module_fq.is_empty() {
755 return None;
756 }
757 let components = python_module_components(file);
758 debug_assert_eq!(
759 module_fq,
760 components.join("."),
761 "module_code_unit must be built from the file's path-derived Python module name"
762 );
763 let structured_fq = python_module_fq_from_components(&components);
764 module_code_unit_from_fq(file, &components, structured_fq)
765}
766
767fn module_code_unit_from_fq(
768 file: &ProjectFile,
769 components: &[String],
770 structured_fq: FqName,
771) -> Option<CodeUnit> {
772 let (short_name, package_components) = components.split_last()?;
773 let package_name = package_components.join(".");
774 Some(CodeUnit::new_fq(
775 file.clone(),
776 CodeUnitType::Module,
777 package_name,
778 short_name.clone(),
779 structured_fq,
780 ))
781}
782
783fn python_class_signature(node: Node<'_>, source: &str) -> String {
784 python_header_with_decorators(node, source)
785}
786
787fn python_function_signature(node: Node<'_>, source: &str) -> String {
788 let header = python_header_with_decorators(node, source);
789 if let Some((head, tail)) = header.rsplit_once('\n') {
790 format!("{head}\n{tail} ...")
791 } else {
792 format!("{header} ...")
793 }
794}
795
796fn python_signature_metadata(signature: String, node: Node<'_>, source: &str) -> SignatureMetadata {
797 let Some(parameters_node) = node.child_by_field_name("parameters") else {
798 return SignatureMetadata::new(signature, Vec::new())
799 .with_dispatch_extensibility(DispatchExtensibility::Open);
800 };
801 let parameter_text = py_node_text(parameters_node, source).trim();
802 let Some(parameters_start) = signature.find(parameter_text) else {
803 return SignatureMetadata::new(signature, Vec::new())
804 .with_dispatch_extensibility(DispatchExtensibility::Open);
805 };
806 let parameters_end = parameters_start + parameter_text.len();
807 let mut search_start = parameters_start;
808 let parameters = python_parameter_label_nodes(parameters_node)
809 .into_iter()
810 .filter_map(|label_node| {
811 let label = py_node_text(label_node, source).trim();
812 if label.is_empty() || search_start > parameters_end {
813 return None;
814 }
815 let haystack = signature.get(search_start..parameters_end)?;
816 let relative_start = haystack.find(label)?;
817 let start_byte = search_start + relative_start;
818 let end_byte = start_byte + label.len();
819 search_start = end_byte;
820 Some(ParameterMetadata::new(label, start_byte, end_byte))
821 })
822 .collect();
823 SignatureMetadata::new(signature, parameters)
824 .with_dispatch_extensibility(DispatchExtensibility::Open)
825}
826
827fn python_parameter_label_nodes(parameters_node: Node<'_>) -> Vec<Node<'_>> {
828 let mut labels = Vec::new();
829 let mut cursor = parameters_node.walk();
830 for child in parameters_node.named_children(&mut cursor) {
831 if let Some(label_node) = python_parameter_label_node(child) {
832 labels.push(label_node);
833 }
834 }
835 labels
836}
837
838pub fn python_parameter_label_node(node: Node<'_>) -> Option<Node<'_>> {
846 match node.kind() {
847 "identifier" => Some(node),
848 "typed_parameter"
849 | "typed_default_parameter"
850 | "default_parameter"
851 | "list_splat_pattern"
852 | "dictionary_splat_pattern"
853 | "keyword_separator" => node.child_by_field_name("name").or_else(|| {
854 let mut cursor = node.walk();
855 node.named_children(&mut cursor)
856 .find_map(python_parameter_label_node)
857 }),
858 _ => None,
859 }
860}
861
862fn python_is_property_mutator(node: Node<'_>, source: &str) -> bool {
863 python_header_with_decorators(node, source)
864 .lines()
865 .map(str::trim)
866 .filter(|line| line.starts_with('@'))
867 .any(|decorator| decorator.ends_with(".setter") || decorator.ends_with(".deleter"))
868}
869
870pub fn python_expanded_comment_start(source: &str, start_byte: usize) -> usize {
871 let line_starts = compute_line_starts(source);
872 let line_index = find_line_index_for_offset(&line_starts, start_byte);
873
874 let mut comment_start = start_byte;
875 for line_idx in (0..line_index).rev() {
876 let line_start = line_starts[line_idx];
877 let line_end = line_starts
878 .get(line_idx + 1)
879 .copied()
880 .unwrap_or(source.len());
881 let line = &source[line_start..line_end];
882 let trimmed = line.trim_start();
883
884 if trimmed.trim().is_empty() {
885 continue;
886 }
887
888 if trimmed.starts_with('#') {
889 comment_start = line_start;
890 continue;
891 }
892
893 break;
894 }
895
896 comment_start
897}
898
899fn python_header_with_decorators(node: Node<'_>, source: &str) -> String {
900 let raw = py_node_text(node, source);
901 let lines: Vec<_> = raw
902 .lines()
903 .map(str::trim_end)
904 .filter(|line| !line.trim().is_empty())
905 .collect();
906 let mut relevant = Vec::new();
907 for line in lines {
908 let trimmed = line.trim_start();
909 if trimmed.starts_with('@')
910 || trimmed.starts_with("def ")
911 || trimmed.starts_with("async def ")
912 || trimmed.starts_with("class ")
913 {
914 relevant.push(trimmed.to_string());
915 if trimmed.starts_with("def ")
916 || trimmed.starts_with("async def ")
917 || trimmed.starts_with("class ")
918 {
919 break;
920 }
921 }
922 }
923 relevant.join("\n")
924}
925
926fn extract_python_supertypes(node: Node<'_>, source: &str) -> Vec<String> {
927 let Some(superclasses) = node.child_by_field_name("superclasses") else {
928 return Vec::new();
929 };
930 let mut result = Vec::new();
931 let mut cursor = superclasses.walk();
932 for child in superclasses.named_children(&mut cursor) {
933 match child.kind() {
934 "identifier" | "attribute" => {
935 let text = py_node_text(child, source).trim();
936 if !text.is_empty() {
937 result.push(text.to_string());
938 }
939 }
940 _ => {}
941 }
942 }
943 result
944}
945
946fn collect_assigned_names(node: Node<'_>, source: &str) -> Vec<String> {
947 let mut names = Vec::new();
948 walk_named_tree_preorder(node, true, |node| {
949 match node.kind() {
950 "attribute" | "subscript" => WalkControl::SkipChildren,
954 "identifier" => {
955 let text = py_node_text(node, source).trim();
956 if !text.is_empty() {
957 names.push(text.to_string());
958 }
959 WalkControl::Continue
960 }
961 _ => WalkControl::Continue,
962 }
963 });
964 names
965}
966
967fn collect_self_assigned_attributes<'tree>(
968 node: Node<'tree>,
969 source: &str,
970 receiver_name: &str,
971) -> Vec<(String, Node<'tree>)> {
972 let mut attributes = Vec::new();
973 collect_direct_self_assigned_attributes(node, source, receiver_name, &mut attributes);
974 attributes
975}
976
977fn collect_direct_self_assigned_attributes<'tree>(
978 node: Node<'tree>,
979 source: &str,
980 receiver_name: &str,
981 attributes: &mut Vec<(String, Node<'tree>)>,
982) {
983 match node.kind() {
984 "attribute" => {
985 let Some(object) = node.child_by_field_name("object") else {
986 return;
987 };
988 if object.kind() != "identifier" || py_node_text(object, source).trim() != receiver_name
989 {
990 return;
991 }
992 let Some(attribute) = node.child_by_field_name("attribute") else {
993 return;
994 };
995 let name = py_node_text(attribute, source).trim();
996 if !name.is_empty() {
997 attributes.push((name.to_string(), attribute));
998 }
999 }
1000 "pattern_list" | "tuple" | "list" | "parenthesized_expression" => {
1001 let mut cursor = node.walk();
1002 for child in node.named_children(&mut cursor) {
1003 collect_direct_self_assigned_attributes(child, source, receiver_name, attributes);
1004 }
1005 }
1006 _ => {}
1007 }
1008}
1009
1010fn python_instance_method_receiver_name(node: Node<'_>, source: &str) -> Option<String> {
1011 if python_function_has_decorator(node, source, "staticmethod")
1012 || python_function_has_decorator(node, source, "classmethod")
1013 {
1014 return None;
1015 }
1016 python_first_parameter_name(node, source)
1017}
1018
1019fn python_function_has_decorator(node: Node<'_>, source: &str, decorator_name: &str) -> bool {
1020 let Some(parent) = node.parent() else {
1021 return false;
1022 };
1023 if parent.kind() != "decorated_definition" {
1024 return false;
1025 }
1026 let mut cursor = parent.walk();
1027 parent
1028 .named_children(&mut cursor)
1029 .filter(|child| child.kind() == "decorator")
1030 .filter_map(|decorator| decorator.named_child(0))
1031 .filter_map(expression_name_node)
1032 .any(|name| py_node_text(name, source).trim() == decorator_name)
1033}
1034
1035fn python_first_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
1036 let parameters = node.child_by_field_name("parameters")?;
1037 let mut cursor = parameters.walk();
1038 parameters
1039 .named_children(&mut cursor)
1040 .find_map(|child| python_parameter_name(child, source))
1041}
1042
1043fn python_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
1044 match node.kind() {
1045 "identifier" => Some(py_node_text(node, source).trim().to_string()),
1046 "typed_parameter"
1047 | "default_parameter"
1048 | "list_splat_pattern"
1049 | "dictionary_splat_pattern" => node
1050 .child_by_field_name("name")
1051 .or_else(|| {
1052 let mut cursor = node.walk();
1053 node.named_children(&mut cursor)
1054 .find(|child| child.kind() == "identifier")
1055 })
1056 .and_then(|name| python_parameter_name(name, source)),
1057 _ => None,
1058 }
1059 .filter(|name| !name.is_empty())
1060}
1061
1062pub fn collect_python_identifiers(node: Node<'_>, source: &str, identifiers: &mut HashSet<String>) {
1063 walk_named_tree_preorder(node, true, |node| {
1064 if node.kind() == "identifier" {
1065 let text = py_node_text(node, source).trim();
1066 if !text.is_empty() {
1067 identifiers.insert(text.to_string());
1068 }
1069 }
1070 WalkControl::Continue
1071 });
1072}
1073
1074pub fn parse_python_tree(source: &str) -> Option<Tree> {
1075 let mut parser = Parser::new();
1076 parser
1077 .set_language(&tree_sitter_python::LANGUAGE.into())
1078 .expect("failed to load python parser");
1079 parser.parse(source, None)
1080}