Skip to main content

brokk_bifrost_python/
declarations.rs

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
15/// Intern one qualified-name segment in the process-global interner.
16fn py_segment(text: &str, kind: SegmentKind) -> SegmentId {
17    segment_interner().intern(text, kind)
18}
19
20/// Build the structured module-path prefix for a Python declaration.
21///
22/// Ordinary modules render as a dotted path such as `mypkg.subpkg.mymodule`,
23/// with each original path component represented by one
24/// [`SegmentKind::Package`] segment. Hidden directories such as `.agent` and
25/// `.github` are also legal components in the analyzer's path-derived Python
26/// convention, but their leading dot is ambiguous after a rendered name has
27/// been joined.
28///
29/// Build the structured name from the file path's original components so
30/// hidden-directory segments stay intact in cold extraction, synthesized module
31/// units, and persisted reconstruction.
32pub 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    /// The structured qualified name matching `path` (M1 dual representation;
101    /// see `.agents/plans/fqname-interned-segments.md`). Tracked independent of
102    /// whether this scope level was actually `capture`d as a `CodeUnit`, so a
103    /// nested class/function that IS captured can always extend an ancestor's
104    /// `fq` even when an intermediate scope level (e.g. a non-captured nested
105    /// function) has no `code_unit` of its own to read `.fq()` from.
106    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        // A nested class (any parent scope, Class or Function) is always joined
283        // with a literal `$` in the legacy convention above, which is exactly
284        // what `SegmentKind::Nested` renders regardless of the preceding
285        // segment's kind; a top-level class has no parent and is a plain `Type`
286        // hanging off the module-path `Package` chain.
287        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        // Mirrors `short_name` above segment-for-segment: a method owned
373        // directly by a class joins with `.` (`Member`), while a function
374        // nested under another function is a local/closure and joins with the
375        // literal `$` that `SegmentKind::Nested` renders.
376        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.imports.push(info);
583        }
584    }
585}
586
587/// Build the [`ParsedFile`] for one Python source file: module unit, type
588/// identifiers, and the declaration walk. `analyzer/python/adapter.rs`'s
589/// `LanguageAdapter::parse_file` is the only caller.
590pub fn parse_python_file(file: &ProjectFile, source: &str, tree: &Tree) -> ParsedFile {
591    let module_fq = python_module_name(file);
592    let mut parsed = ParsedFile::new(module_fq.clone());
593    let root = tree.root_node();
594
595    collect_python_identifiers(root, source, &mut parsed.type_identifiers);
596
597    let module_code_unit = module_code_unit(file, &module_fq);
598    if let Some(module) = module_code_unit.clone() {
599        parsed.add_code_unit(module, root, source, None, None);
600    }
601
602    let overload_decorators = PythonOverloadDecoratorBindings::collect(root, source);
603    let mut visitor = PythonVisitor {
604        file,
605        source,
606        package_name: &module_fq,
607        parsed: &mut parsed,
608        module: module_code_unit,
609        overload_decorators: &overload_decorators,
610    };
611    visitor.visit_container(root, &[], 0);
612
613    parsed
614}
615
616pub fn py_node_text<'a>(node: Node<'_>, source: &'a str) -> &'a str {
617    brokk_bifrost_core::analyzer::common::node_source_text(node, source)
618}
619
620pub fn python_module_name(file: &ProjectFile) -> String {
621    python_module_components(file).join(".")
622}
623
624pub fn module_code_unit(file: &ProjectFile, module_fq: &str) -> Option<CodeUnit> {
625    if module_fq.is_empty() {
626        return None;
627    }
628    let mut components = python_module_components(file);
629    debug_assert_eq!(
630        module_fq,
631        components.join("."),
632        "module_code_unit must be built from the file's path-derived Python module name"
633    );
634    let short_name = components.pop()?;
635    let package_name = components.join(".");
636    Some(CodeUnit::new_fq(
637        file.clone(),
638        CodeUnitType::Module,
639        package_name,
640        short_name,
641        python_module_fq(file),
642    ))
643}
644
645fn python_class_signature(node: Node<'_>, source: &str) -> String {
646    python_header_with_decorators(node, source)
647}
648
649fn python_function_signature(node: Node<'_>, source: &str) -> String {
650    let header = python_header_with_decorators(node, source);
651    if let Some((head, tail)) = header.rsplit_once('\n') {
652        format!("{head}\n{tail} ...")
653    } else {
654        format!("{header} ...")
655    }
656}
657
658fn python_signature_metadata(signature: String, node: Node<'_>, source: &str) -> SignatureMetadata {
659    let Some(parameters_node) = node.child_by_field_name("parameters") else {
660        return SignatureMetadata::new(signature, Vec::new())
661            .with_dispatch_extensibility(DispatchExtensibility::Open);
662    };
663    let parameter_text = py_node_text(parameters_node, source).trim();
664    let Some(parameters_start) = signature.find(parameter_text) else {
665        return SignatureMetadata::new(signature, Vec::new())
666            .with_dispatch_extensibility(DispatchExtensibility::Open);
667    };
668    let parameters_end = parameters_start + parameter_text.len();
669    let mut search_start = parameters_start;
670    let parameters = python_parameter_label_nodes(parameters_node)
671        .into_iter()
672        .filter_map(|label_node| {
673            let label = py_node_text(label_node, source).trim();
674            if label.is_empty() || search_start > parameters_end {
675                return None;
676            }
677            let haystack = signature.get(search_start..parameters_end)?;
678            let relative_start = haystack.find(label)?;
679            let start_byte = search_start + relative_start;
680            let end_byte = start_byte + label.len();
681            search_start = end_byte;
682            Some(ParameterMetadata::new(label, start_byte, end_byte))
683        })
684        .collect();
685    SignatureMetadata::new(signature, parameters)
686        .with_dispatch_extensibility(DispatchExtensibility::Open)
687}
688
689fn python_parameter_label_nodes(parameters_node: Node<'_>) -> Vec<Node<'_>> {
690    let mut labels = Vec::new();
691    let mut cursor = parameters_node.walk();
692    for child in parameters_node.named_children(&mut cursor) {
693        if let Some(label_node) = python_parameter_label_node(child) {
694            labels.push(label_node);
695        }
696    }
697    labels
698}
699
700/// The identifier node that names one parameter's binding.
701///
702/// The grammar gives `default_parameter` and `typed_default_parameter` a
703/// `name` field but gives `typed_parameter` and the two splat patterns none,
704/// so a caller that reads only the field loses the binding name of every
705/// annotated parameter. Every Python surface that names parameters reads them
706/// through this function.
707pub fn python_parameter_label_node(node: Node<'_>) -> Option<Node<'_>> {
708    match node.kind() {
709        "identifier" => Some(node),
710        "typed_parameter"
711        | "typed_default_parameter"
712        | "default_parameter"
713        | "list_splat_pattern"
714        | "dictionary_splat_pattern"
715        | "keyword_separator" => node.child_by_field_name("name").or_else(|| {
716            let mut cursor = node.walk();
717            node.named_children(&mut cursor)
718                .find_map(python_parameter_label_node)
719        }),
720        _ => None,
721    }
722}
723
724fn python_is_property_mutator(node: Node<'_>, source: &str) -> bool {
725    python_header_with_decorators(node, source)
726        .lines()
727        .map(str::trim)
728        .filter(|line| line.starts_with('@'))
729        .any(|decorator| decorator.ends_with(".setter") || decorator.ends_with(".deleter"))
730}
731
732pub fn python_expanded_comment_start(source: &str, start_byte: usize) -> usize {
733    let line_starts = compute_line_starts(source);
734    let line_index = find_line_index_for_offset(&line_starts, start_byte);
735
736    let mut comment_start = start_byte;
737    for line_idx in (0..line_index).rev() {
738        let line_start = line_starts[line_idx];
739        let line_end = line_starts
740            .get(line_idx + 1)
741            .copied()
742            .unwrap_or(source.len());
743        let line = &source[line_start..line_end];
744        let trimmed = line.trim_start();
745
746        if trimmed.trim().is_empty() {
747            continue;
748        }
749
750        if trimmed.starts_with('#') {
751            comment_start = line_start;
752            continue;
753        }
754
755        break;
756    }
757
758    comment_start
759}
760
761fn python_header_with_decorators(node: Node<'_>, source: &str) -> String {
762    let raw = py_node_text(node, source);
763    let lines: Vec<_> = raw
764        .lines()
765        .map(str::trim_end)
766        .filter(|line| !line.trim().is_empty())
767        .collect();
768    let mut relevant = Vec::new();
769    for line in lines {
770        let trimmed = line.trim_start();
771        if trimmed.starts_with('@')
772            || trimmed.starts_with("def ")
773            || trimmed.starts_with("async def ")
774            || trimmed.starts_with("class ")
775        {
776            relevant.push(trimmed.to_string());
777            if trimmed.starts_with("def ")
778                || trimmed.starts_with("async def ")
779                || trimmed.starts_with("class ")
780            {
781                break;
782            }
783        }
784    }
785    relevant.join("\n")
786}
787
788fn extract_python_supertypes(node: Node<'_>, source: &str) -> Vec<String> {
789    let Some(superclasses) = node.child_by_field_name("superclasses") else {
790        return Vec::new();
791    };
792    let mut result = Vec::new();
793    let mut cursor = superclasses.walk();
794    for child in superclasses.named_children(&mut cursor) {
795        match child.kind() {
796            "identifier" | "attribute" => {
797                let text = py_node_text(child, source).trim();
798                if !text.is_empty() {
799                    result.push(text.to_string());
800                }
801            }
802            _ => {}
803        }
804    }
805    result
806}
807
808fn collect_assigned_names(node: Node<'_>, source: &str) -> Vec<String> {
809    let mut names = Vec::new();
810    walk_named_tree_preorder(node, true, |node| {
811        match node.kind() {
812            // An attribute or subscript target (`foo.bar = …`, `foo[i] = …`)
813            // mutates an existing object; it declares neither the receiver nor
814            // the member as a name, so do not descend into it.
815            "attribute" | "subscript" => WalkControl::SkipChildren,
816            "identifier" => {
817                let text = py_node_text(node, source).trim();
818                if !text.is_empty() {
819                    names.push(text.to_string());
820                }
821                WalkControl::Continue
822            }
823            _ => WalkControl::Continue,
824        }
825    });
826    names
827}
828
829fn collect_self_assigned_attributes<'tree>(
830    node: Node<'tree>,
831    source: &str,
832    receiver_name: &str,
833) -> Vec<(String, Node<'tree>)> {
834    let mut attributes = Vec::new();
835    collect_direct_self_assigned_attributes(node, source, receiver_name, &mut attributes);
836    attributes
837}
838
839fn collect_direct_self_assigned_attributes<'tree>(
840    node: Node<'tree>,
841    source: &str,
842    receiver_name: &str,
843    attributes: &mut Vec<(String, Node<'tree>)>,
844) {
845    match node.kind() {
846        "attribute" => {
847            let Some(object) = node.child_by_field_name("object") else {
848                return;
849            };
850            if object.kind() != "identifier" || py_node_text(object, source).trim() != receiver_name
851            {
852                return;
853            }
854            let Some(attribute) = node.child_by_field_name("attribute") else {
855                return;
856            };
857            let name = py_node_text(attribute, source).trim();
858            if !name.is_empty() {
859                attributes.push((name.to_string(), attribute));
860            }
861        }
862        "pattern_list" | "tuple" | "list" | "parenthesized_expression" => {
863            let mut cursor = node.walk();
864            for child in node.named_children(&mut cursor) {
865                collect_direct_self_assigned_attributes(child, source, receiver_name, attributes);
866            }
867        }
868        _ => {}
869    }
870}
871
872fn python_instance_method_receiver_name(node: Node<'_>, source: &str) -> Option<String> {
873    if python_function_has_decorator(node, source, "staticmethod")
874        || python_function_has_decorator(node, source, "classmethod")
875    {
876        return None;
877    }
878    python_first_parameter_name(node, source)
879}
880
881fn python_function_has_decorator(node: Node<'_>, source: &str, decorator_name: &str) -> bool {
882    let Some(parent) = node.parent() else {
883        return false;
884    };
885    if parent.kind() != "decorated_definition" {
886        return false;
887    }
888    let mut cursor = parent.walk();
889    parent
890        .named_children(&mut cursor)
891        .filter(|child| child.kind() == "decorator")
892        .filter_map(|decorator| decorator.named_child(0))
893        .filter_map(expression_name_node)
894        .any(|name| py_node_text(name, source).trim() == decorator_name)
895}
896
897fn python_first_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
898    let parameters = node.child_by_field_name("parameters")?;
899    let mut cursor = parameters.walk();
900    parameters
901        .named_children(&mut cursor)
902        .find_map(|child| python_parameter_name(child, source))
903}
904
905fn python_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
906    match node.kind() {
907        "identifier" => Some(py_node_text(node, source).trim().to_string()),
908        "typed_parameter"
909        | "default_parameter"
910        | "list_splat_pattern"
911        | "dictionary_splat_pattern" => node
912            .child_by_field_name("name")
913            .or_else(|| {
914                let mut cursor = node.walk();
915                node.named_children(&mut cursor)
916                    .find(|child| child.kind() == "identifier")
917            })
918            .and_then(|name| python_parameter_name(name, source)),
919        _ => None,
920    }
921    .filter(|name| !name.is_empty())
922}
923
924pub fn collect_python_identifiers(node: Node<'_>, source: &str, identifiers: &mut HashSet<String>) {
925    walk_named_tree_preorder(node, true, |node| {
926        if node.kind() == "identifier" {
927            let text = py_node_text(node, source).trim();
928            if !text.is_empty() {
929                identifiers.insert(text.to_string());
930            }
931        }
932        WalkControl::Continue
933    });
934}
935
936pub fn parse_python_tree(source: &str) -> Option<Tree> {
937    let mut parser = Parser::new();
938    parser
939        .set_language(&tree_sitter_python::LANGUAGE.into())
940        .expect("failed to load python parser");
941    parser.parse(source, None)
942}