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::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
16/// Intern one qualified-name segment in the process-global interner.
17fn py_segment(text: &str, kind: SegmentKind) -> SegmentId {
18    segment_interner().intern(text, kind)
19}
20
21/// Build the structured module-path prefix for a Python declaration.
22///
23/// Ordinary modules render as a dotted path such as `mypkg.subpkg.mymodule`,
24/// with each original path component represented by one
25/// [`SegmentKind::Package`] segment. Hidden directories such as `.agent` and
26/// `.github` are also legal components in the analyzer's path-derived Python
27/// convention, but their leading dot is ambiguous after a rendered name has
28/// been joined.
29///
30/// Build the structured name from the file path's original components so
31/// hidden-directory segments stay intact in cold extraction, synthesized module
32/// units, and persisted reconstruction.
33pub 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
91/// Find the nearest setuptools import root that contains this source file.
92///
93/// The manifest is parsed as TOML. An unrelated or malformed `pyproject.toml`
94/// does not change the existing `__init__.py` package-root convention.
95fn 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    /// The structured qualified name matching `path` (M1 dual representation;
147    /// see `.agents/plans/fqname-interned-segments.md`). Tracked independent of
148    /// whether this scope level was actually `capture`d as a `CodeUnit`, so a
149    /// nested class/function that IS captured can always extend an ancestor's
150    /// `fq` even when an intermediate scope level (e.g. a non-captured nested
151    /// function) has no `code_unit` of its own to read `.fq()` from.
152    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        // A nested class (any parent scope, Class or Function) is always joined
330        // with a literal `$` in the legacy convention above, which is exactly
331        // what `SegmentKind::Nested` renders regardless of the preceding
332        // segment's kind; a top-level class has no parent and is a plain `Type`
333        // hanging off the module-path `Package` chain.
334        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        // Mirrors `short_name` above segment-for-segment: a method owned
423        // directly by a class joins with `.` (`Member`), while a function
424        // nested under another function is a local/closure and joins with the
425        // literal `$` that `SegmentKind::Nested` renders.
426        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                // Reassigning a class attribute does not mint a new logical
573                // member. Preserve every physical binding range so class-body
574                // references between assignments can select the active one.
575                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
652/// Build the [`ParsedFile`] for one Python source file: module unit, type
653/// identifiers, and the declaration walk. `analyzer/python/adapter.rs`'s
654/// `LanguageAdapter::parse_file` is the only caller.
655pub 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
777/// The identifier node that names one parameter's binding.
778///
779/// The grammar gives `default_parameter` and `typed_default_parameter` a
780/// `name` field but gives `typed_parameter` and the two splat patterns none,
781/// so a caller that reads only the field loses the binding name of every
782/// annotated parameter. Every Python surface that names parameters reads them
783/// through this function.
784pub 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            // An attribute or subscript target (`foo.bar = …`, `foo[i] = …`)
890            // mutates an existing object; it declares neither the receiver nor
891            // the member as a name, so do not descend into it.
892            "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}