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