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.import_statements.push(info.raw_snippet.clone());
583            self.parsed.imports.push(info);
584        }
585    }
586}
587
588/// Build the [`ParsedFile`] for one Python source file: module unit, type
589/// identifiers, and the declaration walk. `analyzer/python/adapter.rs`'s
590/// `LanguageAdapter::parse_file` is the only caller.
591pub 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            // An attribute or subscript target (`foo.bar = …`, `foo[i] = …`)
807            // mutates an existing object; it declares neither the receiver nor
808            // the member as a name, so do not descend into it.
809            "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}