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