Skip to main content

brokk_bifrost_python/graph/
extractor.rs

1//! Python's forward, per-target usage scan: the scoped import closure, the
2//! per-file walk that proves a reference, and the receiver-type facts both this
3//! and the inverted walk resolve through.
4
5use crate::graph::PythonGraphSource;
6use crate::graph::hits::{
7    record_hit, record_import_hit, record_self_receiver_hit, record_unproven_hit,
8};
9use crate::graph::resolver::{
10    annotation_class_qualifier_site, annotation_reference_candidates, member_name,
11    normalized_receiver_type, receiver_annotation_matches_target,
12    resolve_callable_parameter_default_types, resolve_constructor_types, resolve_receiver_type,
13    target_owner_code_unit, top_level_identifier,
14};
15use crate::graph_support::{PythonSource, PythonUsageSource};
16use crate::imports::{PythonImportBinding, parse_python_import_bindings, resolve_fqn_candidates};
17use crate::usage_index::{
18    ModuleBindingEvent, ModuleBindingEventKind, ModuleBindingTimeline, PythonScopeFacts,
19    usage_matching_edges, usage_module_binding_timeline, usage_resolve_module_files,
20    usage_scope_facts,
21};
22use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
23use brokk_bifrost_core::analyzer::usages::local_inference::{
24    LocalBindingsSnapshot, LocalInferenceConfig, LocalInferenceEngine, SymbolResolution,
25};
26use brokk_bifrost_core::analyzer::usages::model::{ImportKind, UsageHit};
27use brokk_bifrost_core::analyzer::usages::{ImportEdge, ImportEdgeKind};
28use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, Language, ProjectFile, Range};
29use brokk_bifrost_core::cancellation::CancellationToken;
30use brokk_bifrost_core::hash::{HashMap, HashSet};
31use brokk_bifrost_core::text_utils::compute_line_starts;
32use rayon::prelude::*;
33use std::collections::BTreeSet;
34use std::sync::{Arc, Mutex};
35use tree_sitter::{Node, Parser, Tree};
36
37pub struct ParsedFile {
38    pub source: Arc<String>,
39    pub tree: Tree,
40}
41
42pub struct PythonProjectGraph {
43    parsed: HashMap<ProjectFile, ParsedFile>,
44}
45
46impl PythonProjectGraph {
47    pub fn scan_files(
48        &self,
49        candidate_files: &HashSet<ProjectFile>,
50        target_file: &ProjectFile,
51    ) -> HashSet<ProjectFile> {
52        candidate_files
53            .iter()
54            .cloned()
55            .chain(std::iter::once(target_file.clone()))
56            .collect()
57    }
58}
59
60pub fn build_python_graph(
61    candidate_files: &HashSet<ProjectFile>,
62    target_file: &ProjectFile,
63    cancellation: Option<&CancellationToken>,
64) -> PythonProjectGraph {
65    let parser_language = tree_sitter_python::LANGUAGE.into();
66    let files: HashSet<ProjectFile> = candidate_files
67        .iter()
68        .cloned()
69        .chain(std::iter::once(target_file.clone()))
70        .collect();
71    let mut parsed = HashMap::default();
72
73    for file in files {
74        if cancellation.is_some_and(CancellationToken::is_cancelled) {
75            break;
76        }
77        let Ok(source) = file.read_to_string() else {
78            continue;
79        };
80        if cancellation.is_some_and(CancellationToken::is_cancelled) {
81            break;
82        }
83        if source.is_empty() {
84            continue;
85        }
86        let mut parser = Parser::new();
87        if parser.set_language(&parser_language).is_err() {
88            continue;
89        }
90        let Some(tree) = parser.parse(source.as_str(), None) else {
91            continue;
92        };
93        if cancellation.is_some_and(CancellationToken::is_cancelled) {
94            break;
95        }
96        parsed.insert(
97            file,
98            ParsedFile {
99                source: Arc::new(source),
100                tree,
101            },
102        );
103    }
104
105    PythonProjectGraph { parsed }
106}
107
108pub fn scan_files_for_seeds(
109    graph: &PythonGraphSource<'_>,
110    python: &dyn PythonUsageSource,
111    project_graph: &PythonProjectGraph,
112    files: &HashSet<ProjectFile>,
113    target: &CodeUnit,
114    seeds: &BTreeSet<(ProjectFile, String)>,
115    cancellation: Option<&CancellationToken>,
116) -> ScanResult {
117    let collected: Mutex<BTreeSet<UsageHit>> = Mutex::new(BTreeSet::new());
118    let unproven_collected: Mutex<BTreeSet<UsageHit>> = Mutex::new(BTreeSet::new());
119    let target_short = top_level_identifier(graph.index, target);
120    let target_member = member_name(graph.index, target);
121    let target_owner = target_owner_code_unit(graph.index, target);
122    // A same-file best-effort for unresolvable receivers is only safe when the
123    // member name is unambiguous in the target's file (exactly one class there
124    // declares it), so `recv.member` can only mean the target.
125    let member_unique_in_target_file = target_member.as_deref().is_some_and(|member| {
126        let owners: HashSet<CodeUnit> = graph
127            .index
128            .declarations(target.source())
129            .into_iter()
130            .filter(|decl| {
131                decl.identifier() == member && target_owner_code_unit(graph.index, decl).is_some()
132            })
133            .filter_map(|decl| target_owner_code_unit(graph.index, &decl))
134            .collect();
135        owners.len() == 1
136    });
137    let files_vec: Vec<&ProjectFile> = files.iter().collect();
138    let parser_language = tree_sitter_python::LANGUAGE.into();
139
140    files_vec.par_iter().for_each(|file| {
141        if cancellation.is_some_and(CancellationToken::is_cancelled) {
142            return;
143        }
144        let owned_source: Option<Arc<String>>;
145        let owned_tree: Option<Tree>;
146        let (source_str, tree_ref) = if let Some(parsed) = project_graph.parsed.get(*file) {
147            (parsed.source.as_str(), &parsed.tree)
148        } else {
149            let Ok(source) = file.read_to_string() else {
150                return;
151            };
152            if source.is_empty() {
153                return;
154            }
155            let mut parser = Parser::new();
156            if parser.set_language(&parser_language).is_err() {
157                return;
158            }
159            let Some(tree) = parser.parse(source.as_str(), None) else {
160                return;
161            };
162            owned_source = Some(Arc::new(source));
163            owned_tree = Some(tree);
164            (
165                owned_source.as_deref().unwrap().as_str(),
166                owned_tree.as_ref().unwrap(),
167            )
168        };
169        if cancellation.is_some_and(CancellationToken::is_cancelled) {
170            return;
171        }
172
173        let edges = {
174            let _scope = brokk_bifrost_core::profiling::scope("python_graph::matching_edges");
175            usage_matching_edges(python, file, seeds)
176        };
177        // This is an AST-name gate, not a source-text resolver. Every usage
178        // accepted by the later structural walk has a matching identifier or
179        // can occur inside an annotation string. Skip files that lack both
180        // before building their scope facts and performing that expensive walk.
181        if !file_may_reference_target(
182            tree_ref.root_node(),
183            source_str,
184            target,
185            target_short.as_str(),
186            target_member.as_deref(),
187            &edges,
188        ) {
189            return;
190        }
191        let raw_module_bindings = {
192            let _scope =
193                brokk_bifrost_core::profiling::scope("python_graph::module_binding_timeline");
194            usage_module_binding_timeline(python, file, || {
195                collect_module_binding_timeline(tree_ref.root_node(), source_str)
196            })
197        };
198        let module_bindings = classify_module_binding_timeline(
199            python,
200            file,
201            raw_module_bindings.as_ref(),
202            seeds,
203            &edges,
204        );
205        let scoped_import_bindings = parse_python_import_bindings(source_str);
206        let target_self_file = *file == target.source();
207        let scope_facts = {
208            let _scope = brokk_bifrost_core::profiling::scope("python_graph::scope_facts");
209            usage_scope_facts(python, file, || {
210                collect_scope_facts_from_parsed_source(
211                    graph,
212                    python,
213                    file,
214                    source_str,
215                    tree_ref.root_node(),
216                )
217            })
218        };
219        let scope_range_index = build_scope_range_index(graph, scope_facts.as_ref());
220
221        let mut local_hits = BTreeSet::new();
222        let mut local_unproven_hits = BTreeSet::new();
223        let line_starts = compute_line_starts(source_str);
224
225        let mut scan_ctx = ScanCtx {
226            python,
227            file,
228            source: source_str,
229            line_starts: &line_starts,
230            graph,
231            target,
232            target_short: &target_short,
233            target_member: target_member.as_deref(),
234            target_owner: target_owner.clone(),
235            target_is_module: target.is_module(),
236            target_source: target.source(),
237            seeds,
238            edges: &edges,
239            target_self_file,
240            member_best_effort_unique: target_self_file && member_unique_in_target_file,
241            raw_module_bindings: raw_module_bindings.as_ref(),
242            module_bindings: &module_bindings,
243            scoped_import_bindings: &scoped_import_bindings,
244            scope_facts: scope_facts.as_ref(),
245            scope_range_index: &scope_range_index,
246            hits: &mut local_hits,
247            unproven_hits: &mut local_unproven_hits,
248        };
249
250        {
251            let _scope = brokk_bifrost_core::profiling::scope("python_graph::scan_tree");
252            scan_node(tree_ref.root_node(), &mut scan_ctx);
253        }
254
255        if !local_hits.is_empty() {
256            let mut sink = collected
257                .lock()
258                .expect("usage hit collector mutex poisoned");
259            sink.extend(local_hits);
260        }
261        if !local_unproven_hits.is_empty() {
262            let mut sink = unproven_collected
263                .lock()
264                .expect("usage unproven hit collector mutex poisoned");
265            sink.extend(local_unproven_hits);
266        }
267    });
268
269    ScanResult {
270        hits: collected
271            .into_inner()
272            .expect("usage hit collector mutex poisoned"),
273        unproven_hits: unproven_collected
274            .into_inner()
275            .expect("usage unproven hit collector mutex poisoned"),
276    }
277}
278
279fn file_may_reference_target(
280    root: Node<'_>,
281    source: &str,
282    target: &CodeUnit,
283    target_short: &str,
284    target_member: Option<&str>,
285    edges: &[ImportEdge],
286) -> bool {
287    if target.is_module() {
288        return true;
289    }
290
291    let mut stack = vec![root];
292    while let Some(node) = stack.pop() {
293        if node.kind() == "string" {
294            // Annotation strings can use an FQN or a re-export alias that is
295            // not an exact identifier node. Keep them for the structured
296            // annotation resolver below.
297            return true;
298        }
299        if node.kind() == "identifier" {
300            let name = slice(node, source);
301            if name == target_short
302                || target_member.is_some_and(|member| name == member)
303                || edges.iter().any(|edge| edge.local_name == name)
304            {
305                return true;
306            }
307        }
308
309        let mut cursor = node.walk();
310        stack.extend(node.named_children(&mut cursor));
311    }
312    false
313}
314
315pub struct ScanResult {
316    pub hits: BTreeSet<UsageHit>,
317    pub unproven_hits: BTreeSet<UsageHit>,
318}
319
320pub struct ScanCtx<'a> {
321    python: &'a dyn PythonUsageSource,
322    pub file: &'a ProjectFile,
323    pub source: &'a str,
324    pub line_starts: &'a [usize],
325    pub graph: &'a PythonGraphSource<'a>,
326    target: &'a CodeUnit,
327    target_short: &'a str,
328    target_member: Option<&'a str>,
329    target_owner: Option<CodeUnit>,
330    target_is_module: bool,
331    target_source: &'a ProjectFile,
332    seeds: &'a BTreeSet<(ProjectFile, String)>,
333    edges: &'a [ImportEdge],
334    target_self_file: bool,
335    /// True when a same-file best-effort is justified for an unresolvable
336    /// receiver: the target is a member, its owner is in this file, and exactly
337    /// one class in this file declares that member name (so `recv.member` with
338    /// an un-inferrable `recv` unambiguously means the target). Cross-file
339    /// untyped receivers stay conservative.
340    member_best_effort_unique: bool,
341    raw_module_bindings: &'a ModuleBindingTimeline,
342    module_bindings: &'a HashMap<String, Vec<ClassifiedModuleBindingEvent>>,
343    scoped_import_bindings: &'a [PythonImportBinding],
344    scope_facts: &'a HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
345    scope_range_index: &'a [ScopeRangeEntry],
346    pub hits: &'a mut BTreeSet<UsageHit>,
347    pub unproven_hits: &'a mut BTreeSet<UsageHit>,
348}
349
350struct ScopeRangeEntry {
351    range: Range,
352    scope: CodeUnit,
353    prefix_max_end: usize,
354}
355
356fn build_scope_range_index(
357    graph: &PythonGraphSource<'_>,
358    scope_facts: &HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
359) -> Vec<ScopeRangeEntry> {
360    let mut entries = scope_facts
361        .keys()
362        .flat_map(|scope| {
363            graph
364                .index
365                .ranges(scope)
366                .into_iter()
367                .map(|range| ScopeRangeEntry {
368                    range,
369                    scope: scope.clone(),
370                    prefix_max_end: 0,
371                })
372        })
373        .collect::<Vec<_>>();
374    entries.sort_by(|left, right| {
375        left.range
376            .start_byte
377            .cmp(&right.range.start_byte)
378            .then_with(|| right.range.end_byte.cmp(&left.range.end_byte))
379            .then_with(|| left.scope.cmp(&right.scope))
380    });
381    let mut max_end = 0;
382    for entry in &mut entries {
383        max_end = max_end.max(entry.range.end_byte);
384        entry.prefix_max_end = max_end;
385    }
386    entries
387}
388
389fn indexed_scope_entry<'entry, 'facts>(
390    scope_range_index: &'entry [ScopeRangeEntry],
391    scope_facts: &'facts HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
392    node: Node<'_>,
393    mut skip_innermost: usize,
394) -> Option<(&'entry CodeUnit, &'facts LocalBindingsSnapshot<String>)> {
395    let mut cursor =
396        scope_range_index.partition_point(|entry| entry.range.start_byte <= node.start_byte());
397    while cursor > 0 {
398        cursor -= 1;
399        let entry = &scope_range_index[cursor];
400        if entry.prefix_max_end < node.end_byte() {
401            return None;
402        }
403        if entry.range.end_byte >= node.end_byte() {
404            if skip_innermost > 0 {
405                skip_innermost -= 1;
406                continue;
407            }
408            return scope_facts
409                .get(&entry.scope)
410                .map(|facts| (&entry.scope, facts));
411        }
412    }
413    None
414}
415
416/// The per-function receiver-type facts enclosing `node`, if any. Shared by the
417/// forward scan ([`ScanCtx`]) and the inverted builder (`PyScan`) so the two
418/// paths resolve a receiver's scope through one place.
419pub fn enclosing_scope_facts<'a>(
420    index: &dyn CodeUnitIndex,
421    file: &ProjectFile,
422    scope_facts: &'a HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
423    node: Node<'_>,
424) -> Option<&'a LocalBindingsSnapshot<String>> {
425    let range = Range {
426        start_byte: node.start_byte(),
427        end_byte: node.end_byte(),
428        start_line: 0,
429        end_line: 0,
430    };
431    let enclosing = index.enclosing_code_unit(file, &range)?;
432    scope_facts.get(&enclosing)
433}
434
435impl ScanCtx<'_> {
436    fn scope_entry_for_node(
437        &self,
438        node: Node<'_>,
439    ) -> Option<(&CodeUnit, &LocalBindingsSnapshot<String>)> {
440        indexed_scope_entry(
441            self.scope_range_index,
442            self.scope_facts,
443            node,
444            usize::from(function_declaration_expression_is_outer_scoped(node)),
445        )
446    }
447
448    fn scope_facts_for_node(&self, node: Node<'_>) -> Option<&LocalBindingsSnapshot<String>> {
449        self.scope_entry_for_node(node).map(|(_, facts)| facts)
450    }
451
452    fn binds_target(&self, ident: &str, node: Node<'_>) -> bool {
453        let scope_entry = self.scope_entry_for_node(node);
454        if self.target_self_file
455            && ident == self.target_short
456            && scope_entry
457                .is_none_or(|(scope, facts)| scope.is_module() || !facts.is_shadowed(ident))
458        {
459            return true;
460        }
461        if scope_entry.is_some_and(|(scope, facts)| !scope.is_module() && facts.is_shadowed(ident))
462        {
463            return false;
464        }
465        self.module_binding_targets_query(ident, node)
466    }
467
468    fn receiver_binds_target(&self, expr: &str, node: Node<'_>) -> bool {
469        if self.binds_target(expr, node) {
470            return true;
471        }
472
473        if self.target_member.is_some() && self.import_edge_visible_for(expr, node) {
474            return true;
475        }
476
477        // `self`/`cls` is implicitly typed as the enclosing class, so a same-file
478        // `self.member` access is a usage of that class's member even though the
479        // receiver is never assigned a type the way a local or parameter is.
480        if matches!(expr, "self" | "cls") && self.self_receiver_matches_target(node) {
481            return true;
482        }
483
484        match enclosing_runtime_parameter_type(expr, node, self.source) {
485            EnclosingParameterType::Typed(raw_type) => {
486                return self.receiver_type_matches_target(&raw_type);
487            }
488            EnclosingParameterType::Untyped => return false,
489            EnclosingParameterType::NotDeclared => {}
490        }
491
492        let Some(scope_facts) = self.scope_facts_for_node(node) else {
493            return false;
494        };
495        let resolution = scope_facts.resolution_for(expr);
496        let Some(raw_type) = resolution
497            .as_precise()
498            .and_then(|targets| targets.iter().next())
499        else {
500            return false;
501        };
502        self.receiver_type_matches_target(raw_type)
503    }
504
505    /// Whether `node` is evaluated in the target member owner's class namespace.
506    /// This includes class-level field initializers and the decorators,
507    /// annotations, and defaults of a method declaration. The method body itself
508    /// executes later with ordinary function scoping, where a bare member name
509    /// does not reach the class namespace.
510    fn node_directly_in_owner_class_body(&self, node: Node<'_>) -> bool {
511        let Some(target_owner) = self.target_owner.as_ref() else {
512            return false;
513        };
514        let range = Range {
515            start_byte: node.start_byte(),
516            end_byte: node.end_byte(),
517            start_line: 0,
518            end_line: 0,
519        };
520        let Some(enclosing) = self.graph.index.enclosing_code_unit(self.file, &range) else {
521            return false;
522        };
523        if &enclosing == target_owner {
524            return true;
525        }
526        if enclosing.is_function() {
527            return target_owner_code_unit(self.graph.index, &enclosing).as_ref()
528                == Some(target_owner)
529                && function_declaration_expression_is_outer_scoped(node);
530        }
531        target_owner_code_unit(self.graph.index, &enclosing).as_ref() == Some(target_owner)
532    }
533
534    /// Whether `expr`'s type is genuinely un-inferrable in `node`'s scope (an
535    /// unseeded receiver such as an unannotated parameter), as opposed to a
536    /// receiver we resolved to some specific — possibly different — type.
537    fn receiver_type_is_unknown(&self, expr: &str, node: Node<'_>) -> bool {
538        match enclosing_runtime_parameter_type(expr, node, self.source) {
539            EnclosingParameterType::Typed(_) => return false,
540            EnclosingParameterType::Untyped => return true,
541            EnclosingParameterType::NotDeclared => {}
542        }
543        match self.scope_facts_for_node(node) {
544            Some(facts) => facts.resolution_for(expr).is_unknown(),
545            None => true,
546        }
547    }
548
549    fn import_edge_visible_for(&self, ident: &str, node: Node<'_>) -> bool {
550        if let Some(scope_facts) = self.scope_facts_for_node(node)
551            && scope_facts.is_shadowed(ident)
552        {
553            return false;
554        }
555        self.module_binding_targets_query(ident, node)
556    }
557
558    fn module_binding_targets_query(&self, ident: &str, node: Node<'_>) -> bool {
559        if let Some(matches) = self.function_import_binding_targets_query(ident, node) {
560            return matches;
561        }
562        self.module_binding_matches_query(ident, node, true, |kind| {
563            kind != ModuleBindingKind::Other
564        })
565    }
566
567    fn module_binding_targets_symbol(&self, ident: &str, node: Node<'_>) -> bool {
568        if let Some(matches) = self.function_import_binding_targets_query(ident, node) {
569            return matches;
570        }
571        let unclassified_named_import = self.edges.iter().any(|edge| {
572            edge.local_name == ident && !matches!(edge.kind, ImportEdgeKind::Namespace)
573        });
574        self.module_binding_matches_query(ident, node, unclassified_named_import, |kind| {
575            kind == ModuleBindingKind::TargetSymbolImport
576        })
577    }
578
579    /// Resolve the nearest function-local import before the module timeline.
580    ///
581    /// Candidate discovery retains all imports. A local import with the same
582    /// binder must still override the module binding only inside its function.
583    fn function_import_binding_targets_query(&self, ident: &str, node: Node<'_>) -> Option<bool> {
584        let binding = self.scoped_import_bindings.iter().rev().find(|binding| {
585            binding.is_function_scoped()
586                && binding.start_byte <= node.start_byte()
587                && binding.scope_start_byte <= node.start_byte()
588                && node.end_byte() <= binding.scope_end_byte
589                && binding.local_name == ident
590        })?;
591        let candidates = resolve_fqn_candidates(self.python, &binding.qualified_name, |name| {
592            self.graph.index.definitions(name).collect()
593        });
594        Some(candidates.iter().any(|candidate| candidate == self.target))
595    }
596
597    fn module_binding_matches_query(
598        &self,
599        ident: &str,
600        node: Node<'_>,
601        unclassified: bool,
602        matches: impl Fn(ModuleBindingKind) -> bool,
603    ) -> bool {
604        if !self.edges.iter().any(|edge| edge.local_name == ident) {
605            return false;
606        }
607        let Some(events) = self.module_bindings.get(ident) else {
608            return unclassified;
609        };
610        let cutoff = if reference_is_deferred_function_body(node) {
611            usize::MAX
612        } else {
613            node.start_byte()
614        };
615        let visible: Vec<_> = events
616            .iter()
617            .filter(|event| event.visible_from <= cutoff)
618            .collect();
619        let start = visible
620            .iter()
621            .rposition(|event| !event.conditional)
622            .unwrap_or(0);
623        visible[start..].iter().any(|event| matches(event.kind))
624    }
625
626    /// Whether the class enclosing `node` is the target member's owner (or a
627    /// subclass of it, for inherited members). Used to resolve `self`/`cls`
628    /// receivers, whose type is the lexically enclosing class.
629    fn self_receiver_matches_target(&self, node: Node<'_>) -> bool {
630        let Some(target_owner) = self.target_owner.as_ref() else {
631            return false;
632        };
633        let range = Range {
634            start_byte: node.start_byte(),
635            end_byte: node.end_byte(),
636            start_line: 0,
637            end_line: 0,
638        };
639        let Some(enclosing) = self.graph.index.enclosing_code_unit(self.file, &range) else {
640            return false;
641        };
642        let enclosing_class = if enclosing.is_class() {
643            enclosing
644        } else {
645            match target_owner_code_unit(self.graph.index, &enclosing) {
646                Some(class) => class,
647                None => return false,
648            }
649        };
650        if &enclosing_class == target_owner {
651            return true;
652        }
653        self.graph
654            .hierarchy
655            .map(|provider| provider.get_ancestors(&enclosing_class))
656            .unwrap_or_default()
657            .into_iter()
658            .any(|ancestor| ancestor == *target_owner)
659    }
660
661    fn receiver_type_matches_target(&self, raw_type: &str) -> bool {
662        if receiver_annotation_matches_target(
663            raw_type,
664            self.edges,
665            self.target_short,
666            self.target_self_file,
667        ) {
668            return true;
669        }
670
671        let Some(target_owner) = self.target_owner.as_ref() else {
672            return false;
673        };
674        let Some(receiver_type) = resolve_receiver_type(
675            self.graph,
676            self.python,
677            self.file,
678            raw_type,
679            self.target_self_file,
680        ) else {
681            return false;
682        };
683        if &receiver_type == target_owner {
684            return true;
685        }
686        self.graph
687            .hierarchy
688            .map(|provider| provider.get_ancestors(&receiver_type))
689            .unwrap_or_default()
690            .into_iter()
691            .any(|ancestor| ancestor == *target_owner)
692    }
693}
694
695pub(crate) fn function_declaration_expression_is_outer_scoped(node: Node<'_>) -> bool {
696    let site_start = node.start_byte();
697    let site_end = node.end_byte();
698    let mut current = node;
699    while let Some(parent) = current.parent() {
700        if parent.kind() == "function_definition" {
701            if parent
702                .child_by_field_name("body")
703                .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
704            {
705                return false;
706            }
707            if parent
708                .child_by_field_name("name")
709                .is_some_and(|name| name.id() == node.id())
710            {
711                return false;
712            }
713            if let Some(parameters) = parent.child_by_field_name("parameters")
714                && parameters.start_byte() <= site_start
715                && site_end <= parameters.end_byte()
716            {
717                let mut parameter = node;
718                while parameter.parent() != Some(parameters) {
719                    let Some(next) = parameter.parent() else {
720                        return false;
721                    };
722                    parameter = next;
723                }
724                let binder = if parameter.kind() == "identifier" {
725                    Some(parameter)
726                } else {
727                    parameter.child_by_field_name("name").or_else(|| {
728                        parameter
729                            .named_child(0)
730                            .filter(|child| child.kind() == "identifier")
731                    })
732                };
733                return binder.is_none_or(|binder| binder.id() != node.id());
734            }
735            return true;
736        }
737        if parent.kind() == "decorated_definition" {
738            return current.kind() == "decorator";
739        }
740        if parent.kind() == "class_definition" {
741            break;
742        }
743        current = parent;
744    }
745    false
746}
747
748fn scan_node(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
749    let mut stack = vec![node];
750    while let Some(node) = stack.pop() {
751        match node.kind() {
752            "import_statement" | "import_from_statement" => {
753                handle_import_candidate(node, ctx);
754                continue;
755            }
756            "identifier" => {
757                if handle_annotation_reference_candidate(node, ctx) {
758                    continue;
759                }
760                handle_identifier_candidate(node, ctx);
761            }
762            "attribute" => {
763                if handle_annotation_reference_candidate(node, ctx) {
764                    continue;
765                }
766                handle_attribute_candidate(node, ctx);
767            }
768            "string_content" => {
769                handle_annotation_reference_candidate(node, ctx);
770            }
771            "keyword_argument" => {
772                handle_keyword_argument_candidate(node, ctx);
773                if let Some(value) = node.child_by_field_name("value") {
774                    stack.push(value);
775                }
776                continue;
777            }
778            _ => {}
779        }
780
781        let mut cursor = node.walk();
782        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
783        children.reverse();
784        stack.extend(children);
785    }
786}
787
788fn handle_annotation_reference_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) -> bool {
789    let Some(candidates) = annotation_reference_candidates(
790        ctx.graph,
791        ctx.python,
792        ctx.file,
793        ctx.source,
794        node,
795        ctx.target_self_file,
796    ) else {
797        return false;
798    };
799
800    if (ctx.target.is_class() || ctx.target.is_field() || ctx.target_member.is_none())
801        && candidates.iter().all(|candidate| *candidate == *ctx.target)
802        && candidates.iter().any(|candidate| *candidate == *ctx.target)
803    {
804        let site = if node.kind() == "attribute" {
805            node.child_by_field_name("attribute").unwrap_or(node)
806        } else {
807            node
808        };
809        record_hit(site, ctx);
810    }
811
812    if let Some(site) = annotation_class_qualifier_site(
813        ctx.graph, ctx.python, ctx.file, ctx.source, node, ctx.target,
814    ) {
815        record_hit(site, ctx);
816    }
817
818    // An attribute annotation the resolver could not turn into a candidate is
819    // not consumed here: the namespace-attribute path still has to see it, and
820    // so do the node's children. A module target also needs that path because
821    // annotation resolution names the terminal declaration, not its qualifier.
822    // `inverted.rs` always falls through for these module qualifier references.
823    if node.kind() == "attribute" && (candidates.is_empty() || ctx.target_is_module) {
824        return false;
825    }
826
827    true
828}
829
830fn handle_keyword_argument_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
831    let (Some(target_member), Some(name), Some(arguments)) = (
832        ctx.target_member,
833        node.child_by_field_name("name"),
834        node.parent(),
835    ) else {
836        return;
837    };
838    if name.kind() != "identifier"
839        || slice(name, ctx.source) != target_member
840        || arguments.kind() != "argument_list"
841    {
842        return;
843    }
844    let Some(call) = arguments.parent().filter(|parent| parent.kind() == "call") else {
845        return;
846    };
847    let Some(function) = call.child_by_field_name("function") else {
848        return;
849    };
850    if function.kind() == "identifier" && slice(function, ctx.source) == "cls" {
851        if ctx.self_receiver_matches_target(function) {
852            record_hit(name, ctx);
853        }
854        return;
855    }
856    let Some(target_owner) = ctx.target_owner.as_ref() else {
857        return;
858    };
859    let scoped_callee_matches = if function.kind() == "identifier" {
860        ctx.scope_facts_for_node(function)
861            .and_then(|facts| {
862                facts
863                    .resolution_for(slice(function, ctx.source))
864                    .as_precise()
865                    .and_then(|targets| targets.iter().next().cloned())
866            })
867            .is_some_and(|raw_type| ctx.receiver_type_matches_target(&raw_type))
868    } else {
869        false
870    };
871    let default_callee_matches = if function.kind() == "identifier" {
872        resolve_callable_parameter_default_types(
873            ctx.graph,
874            ctx.python,
875            ctx.file,
876            ctx.source,
877            function,
878            slice(function, ctx.source),
879        )
880        .into_iter()
881        .any(|class| {
882            &class == target_owner
883                || ctx
884                    .graph
885                    .hierarchy
886                    .map(|provider| provider.get_ancestors(&class))
887                    .unwrap_or_default()
888                    .into_iter()
889                    .any(|ancestor| &ancestor == target_owner)
890        })
891    } else {
892        false
893    };
894    let root_shadowed = leftmost_identifier(function).is_some_and(|root| {
895        ctx.scope_facts_for_node(function)
896            .is_some_and(|facts| facts.is_shadowed(slice(root, ctx.source)))
897    });
898    if root_shadowed && !scoped_callee_matches && !default_callee_matches {
899        return;
900    }
901    let matches = scoped_callee_matches
902        || default_callee_matches
903        || (!root_shadowed
904            && resolve_constructor_types(ctx.graph, ctx.python, ctx.file, ctx.source, function)
905                .into_iter()
906                .any(|class| {
907                    &class == target_owner
908                        || ctx
909                            .graph
910                            .hierarchy
911                            .map(|provider| provider.get_ancestors(&class))
912                            .unwrap_or_default()
913                            .into_iter()
914                            .any(|ancestor| &ancestor == target_owner)
915                }));
916    if matches {
917        record_hit(name, ctx);
918    }
919}
920
921fn leftmost_identifier(mut node: Node<'_>) -> Option<Node<'_>> {
922    loop {
923        match node.kind() {
924            "identifier" => return Some(node),
925            "attribute" => node = node.child_by_field_name("object")?,
926            _ => return None,
927        }
928    }
929}
930
931/// Emit an `Import`-binding hit for `from <mod> import <target>` (the token that
932/// brings the target into this file). Gated on there being an import edge whose
933/// local name is the target — so a same-named symbol imported from a different
934/// module is not falsely counted. Only top-level symbols (not members) are
935/// imported by their own name.
936fn handle_import_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
937    if ctx.target_member.is_some() {
938        return;
939    }
940    if !ctx
941        .edges
942        .iter()
943        .any(|edge| edge.local_name == ctx.target_short)
944    {
945        return;
946    }
947    let mut stack = vec![node];
948    while let Some(node) = stack.pop() {
949        if node.kind() == "identifier" && slice(node, ctx.source) == ctx.target_short {
950            record_import_hit(node, ctx);
951            return;
952        }
953        let mut cursor = node.walk();
954        for child in node.named_children(&mut cursor) {
955            stack.push(child);
956        }
957    }
958}
959
960fn handle_identifier_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
961    if node
962        .parent()
963        .is_some_and(|parent| parent.kind() == "attribute")
964    {
965        return;
966    }
967    let text = slice(node, ctx.source);
968    if text.is_empty() || is_declaration_identifier(node) || decorates_the_target(node, ctx) {
969        return;
970    }
971    if let Some(member) = ctx.target_member {
972        // A constructor call `Owner(...)` invokes the class's `__init__`, so it
973        // is a usage of `__init__` even though `__init__` never appears.
974        if member == "__init__" && is_call_callee(node) && ctx.binds_target(text, node) {
975            record_hit(node, ctx);
976            return;
977        }
978        // For a member target, a bare identifier is a usage only when it names
979        // the member directly in the owner class body (the class namespace).
980        if text == member && ctx.node_directly_in_owner_class_body(node) {
981            record_hit(node, ctx);
982        }
983        return;
984    }
985    if !ctx.binds_target(text, node) {
986        return;
987    }
988    if !ctx.target_is_module
989        && ctx.edges.iter().any(|edge| edge.local_name == text)
990        && !ctx.module_binding_targets_symbol(text, node)
991    {
992        return;
993    }
994    record_hit(node, ctx);
995}
996
997/// Whether `node` sits in a decorator of the very definition the scan is
998/// looking for.
999///
1000/// Python evaluates a decorator expression before the decorated name is bound,
1001/// so `@foo` above `def foo()` names an *outer* `foo` and never the definition
1002/// it decorates. Without this the decorated definition reads as a user of
1003/// itself, which the enclosing-equals-target drop used to hide.
1004fn decorates_the_target(node: Node<'_>, ctx: &ScanCtx<'_>) -> bool {
1005    if ctx.file != ctx.target_source {
1006        return false;
1007    }
1008    let mut current = node;
1009    while let Some(parent) = current.parent() {
1010        if parent.kind() == "decorated_definition" && current.kind() == "decorator" {
1011            let Some(definition) = parent.child_by_field_name("definition") else {
1012                return false;
1013            };
1014            return ctx.graph.index.ranges(ctx.target).iter().any(|range| {
1015                range.start_byte <= definition.start_byte()
1016                    && definition.end_byte() <= range.end_byte
1017            });
1018        }
1019        current = parent;
1020    }
1021    false
1022}
1023
1024fn handle_attribute_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
1025    let Some(object) = node.child_by_field_name("object") else {
1026        return;
1027    };
1028    let Some(attribute) = node.child_by_field_name("attribute") else {
1029        return;
1030    };
1031    let object_text = slice(object, ctx.source);
1032    let attribute_text = slice(attribute, ctx.source);
1033    if let Some(member) = ctx.target_member
1034        && attribute_text == member
1035    {
1036        // A `self.member` / `cls.member` receiver is the current instance / own
1037        // class — a same-owner site, excluded from external usage counts (#1014
1038        // facet B). A typed local of the same type is a different instance and
1039        // stays an external hit.
1040        let is_same_owner_receiver =
1041            matches!(object_text, "self" | "cls") && ctx.self_receiver_matches_target(node);
1042        if is_same_owner_receiver {
1043            record_self_receiver_hit(attribute, ctx);
1044        } else if ctx.receiver_binds_target(object_text, node)
1045            || (object.kind() == "call" && call_result_matches_target(object, ctx))
1046        {
1047            record_hit(attribute, ctx);
1048        } else if member_receiver_match_is_unproven(object, object_text, node, ctx) {
1049            record_unproven_hit(attribute, ctx);
1050        }
1051    }
1052
1053    let object_binds_target = if ctx.target_is_module {
1054        imported_root_targets_module(ctx, object, node)
1055    } else {
1056        ctx.binds_target(object_text, node)
1057    };
1058    if object.kind() == "identifier"
1059        && object_binds_target
1060        && (ctx.target_is_module
1061            || (ctx.target_member.is_none()
1062                && !ctx.edges.iter().any(|edge| {
1063                    matches!(edge.kind, ImportEdgeKind::Namespace) && edge.local_name == object_text
1064                })))
1065    {
1066        record_hit(object, ctx);
1067    }
1068
1069    if ctx.target_is_module
1070        && let Some(module_qualifier) = module_attribute_target_hit(node, ctx)
1071    {
1072        record_hit(module_qualifier, ctx);
1073    }
1074
1075    // A bare member name used as the *object* of an attribute access in the
1076    // owner class body — e.g. the `x` in `@x.setter`/`@x.deleter` decorating a
1077    // property `x` — is a usage of that member.
1078    if let Some(member) = ctx.target_member
1079        && object.kind() == "identifier"
1080        && object_text == member
1081        && ctx.node_directly_in_owner_class_body(object)
1082    {
1083        record_hit(object, ctx);
1084    }
1085
1086    // Best-effort for an un-inferrable receiver: `recv.member` where `recv`'s
1087    // type cannot be resolved is attributed to the target when the target's
1088    // owner is in this file and the member name is unique among local classes
1089    // (so `recv.member` can only mean the target). `self`/`cls` are handled
1090    // structurally above; cross-file untyped receivers stay conservative.
1091    if ctx.member_best_effort_unique
1092        && let Some(member) = ctx.target_member
1093        && attribute_text == member
1094        && object.kind() == "identifier"
1095        && !matches!(object_text, "self" | "cls")
1096        && !ctx.receiver_binds_target(object_text, node)
1097        && ctx.receiver_type_is_unknown(object_text, node)
1098    {
1099        record_hit(attribute, ctx);
1100    }
1101
1102    if let Some(module_binding_target) = module_binding_attribute_target_hit(node, ctx) {
1103        record_hit(module_binding_target, ctx);
1104    }
1105}
1106
1107/// Resolve a top-level symbol written through an imported module binding and
1108/// any number of intermediate modules (`K.feature.DISK`,
1109/// `pkg.core.ops.eye_like`). This includes both `import pkg as K` and
1110/// `from pkg import submodule`; the latter is a named import syntactically but
1111/// still introduces a module binding when the structured module index proves
1112/// that `pkg.submodule` is a workspace module.
1113///
1114/// The written path is assembled only from tree-sitter's `attribute` fields.
1115/// The analyzer's canonical export resolver then proves that the whole path
1116/// names the exact physical target, so re-export aliases remain supported
1117/// without comparing rendered source paths or broadening same-name candidates.
1118fn module_binding_attribute_target_hit<'a>(node: Node<'a>, ctx: &ScanCtx<'_>) -> Option<Node<'a>> {
1119    if ctx.target_member.is_some() {
1120        return None;
1121    }
1122    let (root, attributes) = attribute_chain(node)?;
1123    let terminal = *attributes.last()?;
1124    let terminal_name = slice(terminal, ctx.source);
1125    if terminal_name.is_empty()
1126        || !ctx
1127            .seeds
1128            .iter()
1129            .any(|(_, seed_name)| seed_name == terminal_name)
1130    {
1131        return None;
1132    }
1133
1134    for binding in imported_module_bindings(ctx, root, node) {
1135        let mut written_module = binding.module.clone();
1136        for attribute in attributes
1137            [binding.consumed_attributes.min(attributes.len() - 1)..attributes.len() - 1]
1138            .iter()
1139        {
1140            let segment = slice(*attribute, ctx.source);
1141            if segment.is_empty() {
1142                return None;
1143            }
1144            written_module.push('.');
1145            written_module.push_str(segment);
1146        }
1147        if usage_resolve_module_files(ctx.python, ctx.file, &written_module)
1148            .iter()
1149            .any(|resolved| {
1150                ctx.seeds
1151                    .contains(&(resolved.clone(), terminal_name.to_string()))
1152            })
1153        {
1154            return Some(terminal);
1155        }
1156
1157        let mut written_fqn = binding.module;
1158        for attribute in attributes.iter().skip(binding.consumed_attributes) {
1159            let segment = slice(*attribute, ctx.source);
1160            if segment.is_empty() {
1161                return None;
1162            }
1163            written_fqn.push('.');
1164            written_fqn.push_str(segment);
1165        }
1166        if resolve_fqn_candidates(ctx.python, &written_fqn, |name| {
1167            ctx.graph.index.definitions(name).collect()
1168        })
1169        .into_iter()
1170        .any(|candidate| &candidate == ctx.target)
1171        {
1172            return Some(terminal);
1173        }
1174    }
1175    None
1176}
1177
1178fn imported_root_targets_module(ctx: &ScanCtx<'_>, root: Node<'_>, reference: Node<'_>) -> bool {
1179    imported_module_bindings(ctx, root, reference)
1180        .into_iter()
1181        .any(|binding| {
1182            usage_resolve_module_files(ctx.python, ctx.file, &binding.module)
1183                .into_iter()
1184                .any(|resolved_file| &resolved_file == ctx.target_source)
1185        })
1186}
1187
1188fn module_attribute_target_hit<'a>(node: Node<'a>, ctx: &ScanCtx<'_>) -> Option<Node<'a>> {
1189    let (root, attributes) = attribute_chain(node)?;
1190    if attributes.is_empty() {
1191        return None;
1192    }
1193    for binding in imported_module_bindings(ctx, root, node) {
1194        let mut module_fqn = binding.module;
1195        for attribute in attributes.iter().skip(binding.consumed_attributes) {
1196            let segment = slice(*attribute, ctx.source);
1197            if segment.is_empty() {
1198                return None;
1199            }
1200            if module_fqn.ends_with('.') {
1201                module_fqn.push_str(segment);
1202            } else {
1203                module_fqn.push('.');
1204                module_fqn.push_str(segment);
1205            }
1206            let resolved = usage_resolve_module_files(ctx.python, ctx.file, &module_fqn);
1207            if resolved.is_empty() {
1208                break;
1209            }
1210            if resolved
1211                .iter()
1212                .any(|resolved_file| resolved_file == ctx.target_source)
1213            {
1214                return Some(*attribute);
1215            }
1216        }
1217    }
1218    None
1219}
1220
1221fn call_result_matches_target(call: Node<'_>, ctx: &ScanCtx<'_>) -> bool {
1222    let Some(target_owner) = ctx.target_owner.as_ref() else {
1223        return false;
1224    };
1225    let scope_facts = ctx.scope_facts_for_node(call);
1226    call_result_types(
1227        ctx.graph,
1228        ctx.python,
1229        ctx.file,
1230        ctx.source,
1231        call,
1232        scope_facts,
1233    )
1234    .into_iter()
1235    .any(|class| {
1236        &class == target_owner
1237            || ctx
1238                .graph
1239                .hierarchy
1240                .map(|provider| provider.get_ancestors(&class))
1241                .unwrap_or_default()
1242                .into_iter()
1243                .any(|ancestor| &ancestor == target_owner)
1244    })
1245}
1246
1247pub fn call_result_types(
1248    graph: &PythonGraphSource<'_>,
1249    python: &dyn PythonUsageSource,
1250    file: &ProjectFile,
1251    source: &str,
1252    call: Node<'_>,
1253    scope_facts: Option<&LocalBindingsSnapshot<String>>,
1254) -> Vec<CodeUnit> {
1255    let Some(function) = call.child_by_field_name("function") else {
1256        return Vec::new();
1257    };
1258    let constructed = resolve_constructor_types(graph, python, file, source, function);
1259    if !constructed.is_empty() {
1260        return constructed;
1261    }
1262    let callable_fqns = resolve_callable_fqns(graph, python, file, source, function, scope_facts);
1263    if callable_fqns.is_empty() {
1264        return Vec::new();
1265    }
1266    let callables = callable_fqns
1267        .into_iter()
1268        .flat_map(|callable_fqn| {
1269            resolve_fqn_candidates(python, &callable_fqn, |name| {
1270                graph.index.definitions(name).collect()
1271            })
1272        })
1273        .collect::<Vec<_>>();
1274    let mut classes = Vec::new();
1275    for callable in callables.into_iter().filter(CodeUnit::is_function) {
1276        let Some(raw_type) = callable_return_type_name(graph, python, &callable) else {
1277            continue;
1278        };
1279        if let Some(class) =
1280            resolve_receiver_type(graph, python, callable.source(), &raw_type, true)
1281        {
1282            classes.push(class);
1283        }
1284    }
1285    classes.sort();
1286    classes.dedup();
1287    classes
1288}
1289
1290fn resolve_callable_fqns(
1291    graph: &PythonGraphSource<'_>,
1292    python: &dyn PythonUsageSource,
1293    file: &ProjectFile,
1294    source: &str,
1295    function: Node<'_>,
1296    scope_facts: Option<&LocalBindingsSnapshot<String>>,
1297) -> Vec<String> {
1298    match function.kind() {
1299        "identifier" => {
1300            resolve_identifier_callable_fqns(graph, python, file, source, function, scope_facts)
1301        }
1302        "attribute" => {
1303            resolve_attribute_callable_fqns(graph, python, file, source, function, scope_facts)
1304        }
1305        _ => Vec::new(),
1306    }
1307}
1308
1309fn resolve_identifier_callable_fqns(
1310    graph: &PythonGraphSource<'_>,
1311    python: &dyn PythonUsageSource,
1312    file: &ProjectFile,
1313    source: &str,
1314    function: Node<'_>,
1315    scope_facts: Option<&LocalBindingsSnapshot<String>>,
1316) -> Vec<String> {
1317    let local = slice(function, source);
1318    if local.is_empty() || scope_facts.is_some_and(|facts| facts.is_shadowed(local)) {
1319        return Vec::new();
1320    }
1321    let binder = python.import_binder_of(file);
1322    match binder.bindings.get(local) {
1323        Some(binding) if binding.kind == ImportKind::Named => binding
1324            .imported_name
1325            .as_ref()
1326            .map(|imported| vec![format!("{}.{}", binding.module_specifier, imported)])
1327            .unwrap_or_default(),
1328        _ => graph
1329            .index
1330            .declarations(file)
1331            .into_iter()
1332            .find(|unit| unit.is_function() && unit.identifier() == local)
1333            .map(|unit| vec![unit.fq_name()])
1334            .unwrap_or_default(),
1335    }
1336}
1337
1338fn resolve_attribute_callable_fqns(
1339    graph: &PythonGraphSource<'_>,
1340    python: &dyn PythonUsageSource,
1341    file: &ProjectFile,
1342    source: &str,
1343    function: Node<'_>,
1344    scope_facts: Option<&LocalBindingsSnapshot<String>>,
1345) -> Vec<String> {
1346    let Some(receiver) = function.child_by_field_name("object") else {
1347        return Vec::new();
1348    };
1349    let Some(method) = function.child_by_field_name("attribute") else {
1350        return Vec::new();
1351    };
1352    let method = slice(method, source);
1353    if method.is_empty() {
1354        return Vec::new();
1355    }
1356    let mut fqns = attribute_receiver_classes(graph, python, file, source, receiver, scope_facts)
1357        .into_iter()
1358        .map(|class| format!("{}.{}", class.fq_name(), method))
1359        .collect::<Vec<_>>();
1360    fqns.sort();
1361    fqns.dedup();
1362    fqns
1363}
1364
1365fn attribute_receiver_classes(
1366    graph: &PythonGraphSource<'_>,
1367    python: &dyn PythonUsageSource,
1368    file: &ProjectFile,
1369    source: &str,
1370    receiver: Node<'_>,
1371    scope_facts: Option<&LocalBindingsSnapshot<String>>,
1372) -> Vec<CodeUnit> {
1373    let mut classes = match receiver.kind() {
1374        "identifier" => {
1375            identifier_receiver_classes(graph, python, file, source, receiver, scope_facts)
1376        }
1377        "attribute" => {
1378            if let Some(root) = leftmost_identifier(receiver)
1379                && scope_facts.is_some_and(|facts| facts.is_shadowed(slice(root, source)))
1380            {
1381                Vec::new()
1382            } else {
1383                resolve_constructor_types(graph, python, file, source, receiver)
1384            }
1385        }
1386        _ => Vec::new(),
1387    };
1388    classes.sort();
1389    classes.dedup();
1390    classes
1391}
1392
1393fn identifier_receiver_classes(
1394    graph: &PythonGraphSource<'_>,
1395    python: &dyn PythonUsageSource,
1396    file: &ProjectFile,
1397    source: &str,
1398    receiver: Node<'_>,
1399    scope_facts: Option<&LocalBindingsSnapshot<String>>,
1400) -> Vec<CodeUnit> {
1401    let ident = slice(receiver, source);
1402    if ident.is_empty() {
1403        return Vec::new();
1404    }
1405    if matches!(ident, "self" | "cls")
1406        && let Some(class) = enclosing_class_for_node(graph, file, receiver)
1407    {
1408        return vec![class];
1409    }
1410    if let Some(facts) = scope_facts {
1411        if let Some(raw_type) = facts
1412            .resolution_for(ident)
1413            .as_precise()
1414            .and_then(|targets| targets.iter().next())
1415            && let Some(class) = resolve_receiver_type(graph, python, file, raw_type, false)
1416        {
1417            return vec![class];
1418        }
1419        if facts.is_shadowed(ident) {
1420            return Vec::new();
1421        }
1422    }
1423    resolve_receiver_type(graph, python, file, ident, false)
1424        .into_iter()
1425        .collect()
1426}
1427
1428fn enclosing_class_for_node(
1429    graph: &PythonGraphSource<'_>,
1430    file: &ProjectFile,
1431    node: Node<'_>,
1432) -> Option<CodeUnit> {
1433    let range = Range {
1434        start_byte: node.start_byte(),
1435        end_byte: node.end_byte(),
1436        start_line: 0,
1437        end_line: 0,
1438    };
1439    let enclosing = graph.index.enclosing_code_unit(file, &range)?;
1440    // Python's structural model never separately indexes nested
1441    // function/lambda scopes as their own CodeUnit, so `self`/`cls` can only
1442    // ever need the enclosing unit itself or (a method's owner) exactly one
1443    // `parent_of` hop up — `.take(2)` keeps that bound explicit rather than
1444    // an unbounded walk that could climb past the same-file owner class.
1445    brokk_bifrost_core::analyzer::usages::common::enclosing_owner_chain(enclosing, |unit| {
1446        graph.index.parent_of(unit)
1447    })
1448    .take(2)
1449    .find(|unit| unit.is_class() && unit.source() == file)
1450}
1451
1452fn attribute_chain<'a>(node: Node<'a>) -> Option<(Node<'a>, Vec<Node<'a>>)> {
1453    let mut attributes = Vec::new();
1454    let mut current = node;
1455    loop {
1456        if current.kind() != "attribute" {
1457            return None;
1458        }
1459        attributes.push(current.child_by_field_name("attribute")?);
1460        current = current.child_by_field_name("object")?;
1461        if current.kind() == "identifier" {
1462            attributes.reverse();
1463            return Some((current, attributes));
1464        }
1465    }
1466}
1467
1468struct ImportedModuleBinding {
1469    module: String,
1470    consumed_attributes: usize,
1471}
1472
1473fn imported_module_bindings(
1474    ctx: &ScanCtx<'_>,
1475    root: Node<'_>,
1476    reference: Node<'_>,
1477) -> Vec<ImportedModuleBinding> {
1478    let root_text = slice(root, ctx.source);
1479    if root_text.is_empty() || import_root_shadowed(ctx, root_text, root, reference) {
1480        return Vec::new();
1481    }
1482
1483    if let Some(binding) = ctx.scoped_import_bindings.iter().rev().find(|binding| {
1484        binding.is_function_scoped()
1485            && binding.start_byte <= reference.start_byte()
1486            && binding.scope_start_byte <= reference.start_byte()
1487            && reference.end_byte() <= binding.scope_end_byte
1488            && binding.local_name == root_text
1489    }) {
1490        return if usage_resolve_module_files(ctx.python, ctx.file, &binding.qualified_name)
1491            .is_empty()
1492        {
1493            Vec::new()
1494        } else {
1495            vec![ImportedModuleBinding {
1496                module: binding.qualified_name.clone(),
1497                consumed_attributes: binding.consumed_attributes,
1498            }]
1499        };
1500    }
1501
1502    let Some(events) = ctx.raw_module_bindings.get(root_text) else {
1503        return Vec::new();
1504    };
1505    let cutoff = if reference_is_deferred_function_body(reference) {
1506        usize::MAX
1507    } else {
1508        reference.start_byte()
1509    };
1510    let visible: Vec<_> = events
1511        .iter()
1512        .filter(|event| event.visible_from <= cutoff)
1513        .collect();
1514    let start = visible
1515        .iter()
1516        .rposition(|event| !event.conditional)
1517        .unwrap_or(0);
1518    let mut modules = visible[start..]
1519        .iter()
1520        .filter_map(|event| match &event.kind {
1521            ModuleBindingEventKind::ImportModule {
1522                module,
1523                consumed_attributes,
1524            } => Some(ImportedModuleBinding {
1525                module: module.clone(),
1526                consumed_attributes: *consumed_attributes,
1527            }),
1528            ModuleBindingEventKind::FromImport {
1529                module,
1530                imported_name,
1531            } => {
1532                let submodule = if module.ends_with('.') {
1533                    format!("{module}{imported_name}")
1534                } else {
1535                    format!("{module}.{imported_name}")
1536                };
1537                (!usage_resolve_module_files(ctx.python, ctx.file, &submodule).is_empty())
1538                    .then_some(ImportedModuleBinding {
1539                        module: submodule,
1540                        consumed_attributes: 0,
1541                    })
1542            }
1543            ModuleBindingEventKind::Other => None,
1544        })
1545        .collect::<Vec<_>>();
1546    modules.sort_by(|left, right| {
1547        left.module
1548            .cmp(&right.module)
1549            .then_with(|| left.consumed_attributes.cmp(&right.consumed_attributes))
1550    });
1551    modules.dedup_by(|left, right| {
1552        left.module == right.module && left.consumed_attributes == right.consumed_attributes
1553    });
1554    modules
1555}
1556
1557fn import_root_shadowed(
1558    ctx: &ScanCtx<'_>,
1559    root_text: &str,
1560    root: Node<'_>,
1561    reference: Node<'_>,
1562) -> bool {
1563    ctx.scope_entry_for_node(root)
1564        .or_else(|| ctx.scope_entry_for_node(reference))
1565        .is_some_and(|(scope, facts)| !scope.is_module() && facts.is_shadowed(root_text))
1566        || enclosing_parameters_shadow(root_text, reference, ctx.source)
1567}
1568
1569fn enclosing_parameters_shadow(root_text: &str, reference: Node<'_>, source: &str) -> bool {
1570    let mut current = reference;
1571    while let Some(parent) = current.parent() {
1572        if matches!(parent.kind(), "function_definition" | "lambda") {
1573            let Some(parameters) = parent.child_by_field_name("parameters") else {
1574                return false;
1575            };
1576            let mut cursor = parameters.walk();
1577            return parameters.named_children(&mut cursor).any(|parameter| {
1578                parameter_symbol(parameter, source).as_deref() == Some(root_text)
1579            });
1580        }
1581        current = parent;
1582    }
1583    false
1584}
1585
1586enum EnclosingParameterType {
1587    NotDeclared,
1588    Untyped,
1589    Typed(String),
1590}
1591
1592fn enclosing_runtime_parameter_type(
1593    name: &str,
1594    reference: Node<'_>,
1595    source: &str,
1596) -> EnclosingParameterType {
1597    let site_start = reference.start_byte();
1598    let site_end = reference.end_byte();
1599    let mut current = reference;
1600    while let Some(parent) = current.parent() {
1601        if matches!(parent.kind(), "function_definition" | "lambda")
1602            && parent
1603                .child_by_field_name("body")
1604                .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
1605            && let Some(parameters) = parent.child_by_field_name("parameters")
1606        {
1607            let mut cursor = parameters.walk();
1608            for parameter in parameters.named_children(&mut cursor) {
1609                if parameter_symbol(parameter, source).as_deref() != Some(name) {
1610                    continue;
1611                }
1612                return parameter
1613                    .child_by_field_name("type")
1614                    .and_then(|annotation| normalized_receiver_type(slice(annotation, source)))
1615                    .map_or(EnclosingParameterType::Untyped, |raw_type| {
1616                        EnclosingParameterType::Typed(raw_type)
1617                    });
1618            }
1619        }
1620        current = parent;
1621    }
1622    EnclosingParameterType::NotDeclared
1623}
1624
1625fn member_receiver_match_is_unproven(
1626    object: Node<'_>,
1627    object_text: &str,
1628    node: Node<'_>,
1629    ctx: &ScanCtx<'_>,
1630) -> bool {
1631    if matches!(object_text, "self" | "cls") {
1632        return false;
1633    }
1634    match object.kind() {
1635        "identifier" => {
1636            ctx.receiver_type_is_unknown(object_text, node) && !ctx.member_best_effort_unique
1637        }
1638        "attribute" => true,
1639        _ => false,
1640    }
1641}
1642
1643pub fn slice<'a>(node: Node<'_>, source: &'a str) -> &'a str {
1644    brokk_bifrost_core::analyzer::common::node_source_text(node, source)
1645}
1646
1647/// Whether `node` is the `function` callee of a call expression (`node(...)`).
1648fn is_call_callee(node: Node<'_>) -> bool {
1649    node.parent().is_some_and(|parent| {
1650        parent.kind() == "call"
1651            && parent
1652                .child_by_field_name("function")
1653                .is_some_and(|function| function.id() == node.id())
1654    })
1655}
1656
1657pub fn is_declaration_identifier(node: Node<'_>) -> bool {
1658    let Some(parent) = node.parent() else {
1659        return false;
1660    };
1661    let contains = |container: Node<'_>| {
1662        container.start_byte() <= node.start_byte() && node.end_byte() <= container.end_byte()
1663    };
1664    match parent.kind() {
1665        "class_definition" | "function_definition" => parent
1666            .child_by_field_name("name")
1667            .is_some_and(|name| name.id() == node.id()),
1668        "parameters" | "lambda_parameters" | "list_splat_pattern" | "dictionary_splat_pattern" => {
1669            true
1670        }
1671        "default_parameter" | "typed_parameter" | "typed_default_parameter" => {
1672            parent.child_by_field_name("name").is_some_and(contains)
1673        }
1674        "assignment" | "augmented_assignment" | "for_statement" | "for_in_clause" => {
1675            parent.child_by_field_name("left").is_some_and(contains)
1676        }
1677        "named_expression" => parent.child_by_field_name("name").is_some_and(contains),
1678        "aliased_import" | "import_from_statement" | "import_statement" => true,
1679        _ => false,
1680    }
1681}
1682
1683#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1684enum ModuleBindingKind {
1685    TargetSymbolImport,
1686    TargetModuleImport,
1687    Other,
1688}
1689
1690#[derive(Clone, Copy, Debug)]
1691struct ClassifiedModuleBindingEvent {
1692    visible_from: usize,
1693    conditional: bool,
1694    kind: ModuleBindingKind,
1695}
1696
1697pub fn collect_module_binding_timeline(root: Node<'_>, source: &str) -> ModuleBindingTimeline {
1698    let mut timeline = ModuleBindingTimeline::default();
1699    let mut stack = vec![root];
1700    while let Some(node) = stack.pop() {
1701        match node.kind() {
1702            "function_definition" | "class_definition" => {
1703                if let Some(name) = node.child_by_field_name("name") {
1704                    record_module_binding(
1705                        &mut timeline,
1706                        slice(name, source),
1707                        node.end_byte(),
1708                        binding_is_conditional(node),
1709                        ModuleBindingEventKind::Other,
1710                    );
1711                }
1712                continue;
1713            }
1714            "import_statement" | "import_from_statement" => {
1715                collect_import_binding_events(node, source, &mut timeline);
1716                continue;
1717            }
1718            "assignment" | "augmented_assignment" | "named_expression" => {
1719                if let Some(left) = node.child_by_field_name("left") {
1720                    record_local_binding_targets(
1721                        left,
1722                        source,
1723                        node.end_byte(),
1724                        binding_is_conditional(node),
1725                        &mut timeline,
1726                    );
1727                }
1728                continue;
1729            }
1730            "for_statement" => {
1731                if let Some(left) = node.child_by_field_name("left") {
1732                    record_local_binding_targets(
1733                        left,
1734                        source,
1735                        left.end_byte(),
1736                        true,
1737                        &mut timeline,
1738                    );
1739                }
1740            }
1741            _ => {}
1742        }
1743
1744        let mut cursor = node.walk();
1745        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
1746        children.reverse();
1747        stack.extend(children);
1748    }
1749    for events in timeline.values_mut() {
1750        events.sort_by_key(|event| event.visible_from);
1751    }
1752    timeline
1753}
1754
1755fn collect_import_binding_events(
1756    node: Node<'_>,
1757    source: &str,
1758    timeline: &mut ModuleBindingTimeline,
1759) {
1760    if node.kind() == "import_statement" {
1761        let mut cursor = node.walk();
1762        for imported in node.children_by_field_name("name", &mut cursor) {
1763            let name = imported.child_by_field_name("name").unwrap_or(imported);
1764            let Some(local) = imported
1765                .child_by_field_name("alias")
1766                .or_else(|| first_identifier(name))
1767            else {
1768                continue;
1769            };
1770            let module = slice(name, source).trim();
1771            let consumed_attributes = if imported.child_by_field_name("alias").is_some() {
1772                0
1773            } else {
1774                parse_symbol_path(Language::Python, module)
1775                    .len()
1776                    .saturating_sub(1)
1777            };
1778            record_module_binding(
1779                timeline,
1780                slice(local, source),
1781                node.end_byte(),
1782                binding_is_conditional(node),
1783                ModuleBindingEventKind::ImportModule {
1784                    module: module.to_string(),
1785                    consumed_attributes,
1786                },
1787            );
1788        }
1789        return;
1790    }
1791
1792    let Some(module_node) = node.child_by_field_name("module_name") else {
1793        return;
1794    };
1795    let module = slice(module_node, source).trim();
1796    let mut cursor = node.walk();
1797    for imported in node.children_by_field_name("name", &mut cursor) {
1798        if imported.kind() == "wildcard_import" {
1799            continue;
1800        }
1801        let name = imported.child_by_field_name("name").unwrap_or(imported);
1802        let Some(imported_identifier) = last_identifier(name) else {
1803            continue;
1804        };
1805        let imported_name = slice(imported_identifier, source).trim();
1806        let Some(local) = imported
1807            .child_by_field_name("alias")
1808            .or_else(|| last_identifier(name))
1809        else {
1810            continue;
1811        };
1812        record_module_binding(
1813            timeline,
1814            slice(local, source),
1815            node.end_byte(),
1816            binding_is_conditional(node),
1817            ModuleBindingEventKind::FromImport {
1818                module: module.to_string(),
1819                imported_name: imported_name.to_string(),
1820            },
1821        );
1822    }
1823}
1824
1825fn classify_module_binding_timeline(
1826    python: &dyn PythonUsageSource,
1827    file: &ProjectFile,
1828    timeline: &ModuleBindingTimeline,
1829    seeds: &BTreeSet<(ProjectFile, String)>,
1830    edges: &[ImportEdge],
1831) -> HashMap<String, Vec<ClassifiedModuleBindingEvent>> {
1832    let mut classified = HashMap::default();
1833    let mut module_targets: HashMap<String, bool> = HashMap::default();
1834    let relevant_locals: HashSet<&str> =
1835        edges.iter().map(|edge| edge.local_name.as_str()).collect();
1836    for (local, events) in timeline {
1837        if !relevant_locals.contains(local.as_str()) {
1838            continue;
1839        }
1840        let classified_events = events
1841            .iter()
1842            .map(|event| {
1843                let kind = match &event.kind {
1844                    ModuleBindingEventKind::ImportModule { module, .. } => {
1845                        if *module_targets
1846                            .entry(module.clone())
1847                            .or_insert_with(|| module_contains_seed(python, file, module, seeds))
1848                        {
1849                            ModuleBindingKind::TargetModuleImport
1850                        } else {
1851                            ModuleBindingKind::Other
1852                        }
1853                    }
1854                    ModuleBindingEventKind::FromImport {
1855                        module,
1856                        imported_name,
1857                    } => {
1858                        let direct = usage_resolve_module_files(python, file, module).iter().any(
1859                            |resolved| seeds.contains(&(resolved.clone(), imported_name.clone())),
1860                        );
1861                        let submodule = if module.ends_with('.') {
1862                            format!("{module}{imported_name}")
1863                        } else {
1864                            format!("{module}.{imported_name}")
1865                        };
1866                        let imports_target_module =
1867                            *module_targets.entry(submodule.clone()).or_insert_with(|| {
1868                                module_contains_seed(python, file, &submodule, seeds)
1869                            });
1870                        if direct {
1871                            ModuleBindingKind::TargetSymbolImport
1872                        } else if imports_target_module {
1873                            ModuleBindingKind::TargetModuleImport
1874                        } else {
1875                            ModuleBindingKind::Other
1876                        }
1877                    }
1878                    ModuleBindingEventKind::Other => ModuleBindingKind::Other,
1879                };
1880                ClassifiedModuleBindingEvent {
1881                    visible_from: event.visible_from,
1882                    conditional: event.conditional,
1883                    kind,
1884                }
1885            })
1886            .collect();
1887        classified.insert(local.clone(), classified_events);
1888    }
1889    classified
1890}
1891
1892fn module_contains_seed(
1893    python: &dyn PythonUsageSource,
1894    file: &ProjectFile,
1895    module: &str,
1896    seeds: &BTreeSet<(ProjectFile, String)>,
1897) -> bool {
1898    usage_resolve_module_files(python, file, module)
1899        .iter()
1900        .any(|resolved| seeds.iter().any(|(seed_file, _)| seed_file == resolved))
1901}
1902
1903fn record_module_binding(
1904    timeline: &mut ModuleBindingTimeline,
1905    name: &str,
1906    visible_from: usize,
1907    conditional: bool,
1908    kind: ModuleBindingEventKind,
1909) {
1910    let name = name.trim();
1911    if name.is_empty() {
1912        return;
1913    }
1914    timeline
1915        .entry(name.to_string())
1916        .or_default()
1917        .push(ModuleBindingEvent {
1918            visible_from,
1919            conditional,
1920            kind,
1921        });
1922}
1923
1924fn record_local_binding_targets(
1925    target: Node<'_>,
1926    source: &str,
1927    visible_from: usize,
1928    conditional: bool,
1929    timeline: &mut ModuleBindingTimeline,
1930) {
1931    let mut stack = vec![target];
1932    while let Some(node) = stack.pop() {
1933        if node.kind() == "identifier" {
1934            record_module_binding(
1935                timeline,
1936                slice(node, source),
1937                visible_from,
1938                conditional,
1939                ModuleBindingEventKind::Other,
1940            );
1941            continue;
1942        }
1943        if matches!(node.kind(), "attribute" | "subscript") {
1944            continue;
1945        }
1946        let mut cursor = node.walk();
1947        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
1948        children.reverse();
1949        stack.extend(children);
1950    }
1951}
1952
1953fn binding_is_conditional(mut node: Node<'_>) -> bool {
1954    while let Some(parent) = node.parent() {
1955        if matches!(
1956            parent.kind(),
1957            "if_statement"
1958                | "try_statement"
1959                | "except_clause"
1960                | "match_statement"
1961                | "case_clause"
1962                | "for_statement"
1963                | "while_statement"
1964        ) {
1965            return true;
1966        }
1967        if matches!(
1968            parent.kind(),
1969            "module" | "function_definition" | "class_definition"
1970        ) {
1971            return false;
1972        }
1973        node = parent;
1974    }
1975    false
1976}
1977
1978fn first_identifier(node: Node<'_>) -> Option<Node<'_>> {
1979    identifier_extreme(node, false)
1980}
1981
1982fn last_identifier(node: Node<'_>) -> Option<Node<'_>> {
1983    identifier_extreme(node, true)
1984}
1985
1986fn identifier_extreme(node: Node<'_>, last: bool) -> Option<Node<'_>> {
1987    let mut best = None;
1988    let mut stack = vec![node];
1989    while let Some(node) = stack.pop() {
1990        if node.kind() == "identifier" {
1991            if best.is_none_or(|current: Node<'_>| {
1992                if last {
1993                    node.start_byte() > current.start_byte()
1994                } else {
1995                    node.start_byte() < current.start_byte()
1996                }
1997            }) {
1998                best = Some(node);
1999            }
2000            continue;
2001        }
2002        let mut cursor = node.walk();
2003        stack.extend(node.named_children(&mut cursor));
2004    }
2005    best
2006}
2007
2008fn reference_is_deferred_function_body(node: Node<'_>) -> bool {
2009    let site_start = node.start_byte();
2010    let site_end = node.end_byte();
2011    let mut current = node;
2012    while let Some(parent) = current.parent() {
2013        if matches!(parent.kind(), "function_definition" | "lambda")
2014            && parent
2015                .child_by_field_name("body")
2016                .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
2017        {
2018            return true;
2019        }
2020        current = parent;
2021    }
2022    false
2023}
2024
2025pub fn collect_assigned_identifiers(node: Node<'_>, source: &str, out: &mut HashSet<String>) {
2026    let mut stack = vec![node];
2027    while let Some(node) = stack.pop() {
2028        if node.kind() == "identifier" {
2029            let text = slice(node, source).trim();
2030            if !text.is_empty() {
2031                out.insert(text.to_string());
2032            }
2033            continue;
2034        }
2035
2036        let mut cursor = node.walk();
2037        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2038        children.reverse();
2039        stack.extend(children);
2040    }
2041}
2042
2043pub fn collect_scope_facts_from_parsed_source(
2044    graph: &PythonGraphSource<'_>,
2045    python: &dyn PythonUsageSource,
2046    file: &ProjectFile,
2047    source: &str,
2048    root: Node<'_>,
2049) -> PythonScopeFacts {
2050    let mut factory_return_types = collect_factory_return_types_from_root(root, source);
2051    collect_imported_factory_return_types(graph, python, file, &mut factory_return_types);
2052    collect_scope_facts_with_factory_returns(graph, file, source, &factory_return_types)
2053}
2054
2055fn collect_imported_factory_return_types(
2056    graph: &PythonGraphSource<'_>,
2057    python: &dyn PythonUsageSource,
2058    file: &ProjectFile,
2059    factory_return_types: &mut HashMap<String, String>,
2060) {
2061    let binder = python.import_binder_of(file);
2062    for (local, binding) in &binder.bindings {
2063        if !matches!(binding.kind, ImportKind::Named) {
2064            continue;
2065        }
2066        let Some(imported) = binding.imported_name.as_deref() else {
2067            continue;
2068        };
2069        let fqn = format!("{}.{}", binding.module_specifier, imported);
2070        let units =
2071            resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
2072        for unit in units {
2073            if unit.is_function() {
2074                if let Some(return_type) = callable_return_type_name(graph, python, &unit) {
2075                    factory_return_types
2076                        .entry(local.clone())
2077                        .or_insert(return_type);
2078                }
2079                continue;
2080            }
2081            if !unit.is_class() {
2082                continue;
2083            }
2084            factory_return_types
2085                .entry(local.clone())
2086                .or_insert_with(|| unit.identifier().to_string());
2087            collect_imported_class_method_return_types(
2088                graph,
2089                python,
2090                local,
2091                &unit,
2092                factory_return_types,
2093            );
2094        }
2095    }
2096}
2097
2098fn collect_imported_class_method_return_types(
2099    graph: &PythonGraphSource<'_>,
2100    python: &dyn PythonSource,
2101    local_class_name: &str,
2102    class_unit: &CodeUnit,
2103    factory_return_types: &mut HashMap<String, String>,
2104) {
2105    for member in graph.index.direct_children(class_unit) {
2106        if !member.is_function() {
2107            continue;
2108        }
2109        let Some(return_type) = callable_return_type_name(graph, python, &member) else {
2110            continue;
2111        };
2112        factory_return_types
2113            .entry(format!("{}.{}", local_class_name, member.identifier()))
2114            .or_insert(return_type);
2115    }
2116}
2117
2118fn callable_return_type_name(
2119    graph: &PythonGraphSource<'_>,
2120    python: &dyn PythonSource,
2121    callable: &CodeUnit,
2122) -> Option<String> {
2123    // The analyzer's already-parsed whole-file tree when it has one. This runs
2124    // once per imported class member, and the fallback below clones the entire
2125    // file source and builds a fresh `Parser` per declaration range. Same
2126    // prepared-syntax fast path C++ resolution uses for the same reason.
2127    if let Some(prepared) = python.prepared_syntax(callable.source()) {
2128        #[cfg(any(test, feature = "test-support"))]
2129        note_callable_return_type_lookup_for_test(true);
2130        return callable_return_type_name_in_tree(
2131            graph,
2132            callable,
2133            prepared.source(),
2134            prepared.tree().root_node(),
2135        );
2136    }
2137    #[cfg(any(test, feature = "test-support"))]
2138    note_callable_return_type_lookup_for_test(false);
2139    let source = graph.index.indexed_source(callable.source())?;
2140    declaration_source_slices(graph, callable, &source)
2141        .into_iter()
2142        .find_map(|declaration_source| {
2143            let mut parser = Parser::new();
2144            parser
2145                .set_language(&tree_sitter_python::LANGUAGE.into())
2146                .ok()?;
2147            let tree = parser.parse(declaration_source, None)?;
2148            let function = first_function_definition(tree.root_node())?;
2149            factory_return_type(function, declaration_source)
2150        })
2151}
2152
2153/// How `callable_return_type_name` answered, counted per arm.
2154///
2155/// Both arms are counted, not just the slow one: a test that asserted only
2156/// "no reparses" would pass vacuously on a fixture that never reaches this
2157/// function at all.
2158#[cfg(any(test, feature = "test-support"))]
2159#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
2160pub struct CallableReturnTypeLookupCounts {
2161    /// Answered from the analyzer's already-parsed whole-file tree.
2162    pub prepared: usize,
2163    /// Fell back to cloning the file source and building a fresh parser.
2164    pub reparsed: usize,
2165}
2166
2167#[cfg(any(test, feature = "test-support"))]
2168thread_local! {
2169    static CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST: std::cell::Cell<CallableReturnTypeLookupCounts> =
2170        const { std::cell::Cell::new(CallableReturnTypeLookupCounts { prepared: 0, reparsed: 0 }) };
2171}
2172
2173#[cfg(any(test, feature = "test-support"))]
2174fn note_callable_return_type_lookup_for_test(from_prepared_syntax: bool) {
2175    CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST.with(|counts| {
2176        let mut observed = counts.get();
2177        if from_prepared_syntax {
2178            observed.prepared += 1;
2179        } else {
2180            observed.reparsed += 1;
2181        }
2182        counts.set(observed);
2183    });
2184}
2185
2186/// Runs `body` and reports which arm each `callable_return_type_name` call
2187/// took. An analyzer that holds prepared syntax must report zero reparses.
2188#[cfg(any(test, feature = "test-support"))]
2189pub fn with_callable_return_type_lookup_counter_for_test<T>(
2190    body: impl FnOnce() -> T,
2191) -> (T, CallableReturnTypeLookupCounts) {
2192    CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST.with(|counts| {
2193        counts.set(CallableReturnTypeLookupCounts::default());
2194        let result = body();
2195        let observed = counts.get();
2196        counts.set(CallableReturnTypeLookupCounts::default());
2197        (result, observed)
2198    })
2199}
2200
2201/// The declaration walk both arms of [`callable_return_type_name`] share.
2202///
2203/// `source` and `root` must be the same snapshot: `factory_return_type` slices
2204/// `source` at the node's byte offsets, so a whole-file tree needs whole-file
2205/// source and a per-range reparse needs that range's slice.
2206fn callable_return_type_name_in_tree(
2207    graph: &PythonGraphSource<'_>,
2208    callable: &CodeUnit,
2209    source: &str,
2210    root: Node<'_>,
2211) -> Option<String> {
2212    let mut ranges = graph.index.ranges(callable);
2213    ranges.sort_by_key(|range| range.start_byte);
2214    ranges.into_iter().find_map(|range| {
2215        let declaration = root.descendant_for_byte_range(range.start_byte, range.end_byte)?;
2216        let function = first_function_definition(declaration)?;
2217        factory_return_type(function, source)
2218    })
2219}
2220
2221fn first_function_definition(root: Node<'_>) -> Option<Node<'_>> {
2222    let mut stack = vec![root];
2223    while let Some(node) = stack.pop() {
2224        if node.kind() == "function_definition" {
2225            return Some(node);
2226        }
2227        let mut cursor = node.walk();
2228        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2229        children.reverse();
2230        stack.extend(children);
2231    }
2232    None
2233}
2234
2235fn collect_scope_facts_with_factory_returns(
2236    graph: &PythonGraphSource<'_>,
2237    file: &ProjectFile,
2238    source: &str,
2239    factory_return_types: &HashMap<String, String>,
2240) -> PythonScopeFacts {
2241    let declarations = graph.index.declarations(file);
2242    let mut class_facts_by_name: HashMap<String, LocalBindingsSnapshot<String>> =
2243        HashMap::default();
2244    for declaration in declarations
2245        .iter()
2246        .filter(|declaration| declaration.is_class())
2247    {
2248        let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2249            continue;
2250        };
2251        let facts = collect_scope_facts_from_source(
2252            &declaration_source,
2253            ScopeFactTraversal::Class,
2254            true,
2255            Some(declaration.short_name()),
2256            factory_return_types,
2257        );
2258        class_facts_by_name.insert(
2259            declaration.short_name().to_string(),
2260            facts.filtered_visible_bindings(|symbol, _| symbol.starts_with("self.")),
2261        );
2262    }
2263
2264    let mut scope_facts = HashMap::default();
2265    for declaration in declarations
2266        .iter()
2267        .filter(|declaration| declaration.is_function())
2268    {
2269        let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2270            continue;
2271        };
2272        // fqname-M4: package-less short_name owner, matched below against
2273        // `class_facts_by_name` keys built from `short_name()`; `fq.parent()`
2274        // (`default_parent_fq_name`) would render the package-qualified owner,
2275        // a different string that would never hit in that map.
2276        let owner = declaration
2277            .short_name()
2278            .rsplit_once('.')
2279            .map(|(owner, _)| owner);
2280        let mut facts = collect_scope_facts_from_source(
2281            &declaration_source,
2282            ScopeFactTraversal::Function,
2283            false,
2284            owner,
2285            factory_return_types,
2286        );
2287        if let Some(owner) = owner
2288            && let Some(class_facts) = class_facts_by_name.get(owner)
2289        {
2290            facts = facts.merged_with_visible(class_facts);
2291        }
2292        scope_facts.insert(declaration.clone(), facts);
2293    }
2294
2295    // Module-level statements (e.g. a top-level `f = Foo()`) form their own scope.
2296    // `enclosing_code_unit` resolves a top-level usage to the module CodeUnit, so
2297    // its bindings must be recorded too, otherwise constructed-local receivers
2298    // used at module scope resolve to no type.
2299    for declaration in declarations.iter().filter(|d| d.is_module()) {
2300        let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2301            continue;
2302        };
2303        let facts = collect_scope_facts_from_source(
2304            &declaration_source,
2305            ScopeFactTraversal::Module,
2306            false,
2307            None,
2308            factory_return_types,
2309        );
2310        scope_facts.insert(declaration.clone(), facts);
2311    }
2312    scope_facts
2313}
2314
2315fn declaration_source(
2316    graph: &PythonGraphSource<'_>,
2317    declaration: &CodeUnit,
2318    file_source: &str,
2319) -> Option<String> {
2320    let slices = declaration_source_slices(graph, declaration, file_source);
2321    (!slices.is_empty()).then(|| slices.join("\n\n"))
2322}
2323
2324fn declaration_source_slices<'a>(
2325    graph: &PythonGraphSource<'_>,
2326    declaration: &CodeUnit,
2327    file_source: &'a str,
2328) -> Vec<&'a str> {
2329    let mut ranges = graph.index.ranges(declaration);
2330    ranges.sort_by_key(|range| range.start_byte);
2331    ranges
2332        .into_iter()
2333        .filter_map(|range| file_source.get(range.start_byte..range.end_byte))
2334        .collect()
2335}
2336
2337fn collect_scope_facts_from_source(
2338    source: &str,
2339    traversal: ScopeFactTraversal,
2340    allow_self_receivers: bool,
2341    current_class: Option<&str>,
2342    factory_return_types: &HashMap<String, String>,
2343) -> LocalBindingsSnapshot<String> {
2344    let events = collect_scope_fact_events(source, traversal);
2345    collect_scope_facts_from_events(
2346        &events,
2347        allow_self_receivers,
2348        current_class,
2349        factory_return_types,
2350    )
2351}
2352
2353pub fn collect_function_scope_facts_from_node(
2354    function: Node<'_>,
2355    source: &str,
2356) -> LocalBindingsSnapshot<String> {
2357    let mut events = Vec::new();
2358    if function.kind() == "lambda" {
2359        if let Some(parameters) = function.child_by_field_name("parameters") {
2360            collect_parameter_events(parameters, source, &mut events);
2361        }
2362    } else {
2363        collect_scope_fact_events_from_node(
2364            function,
2365            source,
2366            ScopeFactTraversal::Function,
2367            &mut events,
2368        );
2369    }
2370    collect_scope_facts_from_events(&events, false, None, &HashMap::default())
2371}
2372
2373fn collect_scope_facts_from_events(
2374    events: &[ScopeFactEvent],
2375    allow_self_receivers: bool,
2376    current_class: Option<&str>,
2377    factory_return_types: &HashMap<String, String>,
2378) -> LocalBindingsSnapshot<String> {
2379    let mut engine = LocalInferenceEngine::new(LocalInferenceConfig::default());
2380    let globals: HashSet<&str> = events
2381        .iter()
2382        .filter_map(|event| match event {
2383            ScopeFactEvent::Global { symbol } => Some(symbol.as_str()),
2384            _ => None,
2385        })
2386        .collect();
2387    let nonlocals: HashSet<&str> = events
2388        .iter()
2389        .filter_map(|event| match event {
2390            ScopeFactEvent::Nonlocal { symbol } => Some(symbol.as_str()),
2391            _ => None,
2392        })
2393        .collect();
2394    for symbol in &nonlocals {
2395        engine.declare_shadow((*symbol).to_string());
2396    }
2397    for event in events {
2398        if let ScopeFactEvent::Parameter { symbol, .. } = event
2399            && !globals.contains(symbol.as_str())
2400            && !nonlocals.contains(symbol.as_str())
2401            && !engine.is_shadowed(symbol)
2402        {
2403            engine.declare_shadow(symbol.clone());
2404        }
2405    }
2406
2407    let mut changed = true;
2408    while changed {
2409        changed = false;
2410        let mut aliases = Vec::new();
2411        for event in events {
2412            match event {
2413                ScopeFactEvent::Parameter {
2414                    symbol,
2415                    annotation: Some(annotation),
2416                }
2417                | ScopeFactEvent::Annotation { symbol, annotation } => {
2418                    if globals.contains(symbol.as_str()) || nonlocals.contains(symbol.as_str()) {
2419                        continue;
2420                    }
2421                    apply_annotation_event(
2422                        symbol,
2423                        annotation,
2424                        allow_self_receivers,
2425                        &mut engine,
2426                        &mut changed,
2427                    );
2428                }
2429                ScopeFactEvent::Parameter {
2430                    annotation: None, ..
2431                } => {}
2432                ScopeFactEvent::Assignment { lhs, rhs } => {
2433                    if globals.contains(lhs.as_str()) {
2434                        continue;
2435                    }
2436                    if !engine.is_shadowed(lhs) {
2437                        engine.declare_shadow(lhs.clone());
2438                    }
2439                    if lhs.starts_with("self.") && !allow_self_receivers {
2440                        continue;
2441                    }
2442
2443                    match rhs {
2444                        AssignmentRhs::Call(callee) => {
2445                            if !engine.is_shadowed(callee) {
2446                                if let Some(receiver_type) = factory_return_type_for_callee(
2447                                    callee,
2448                                    current_class,
2449                                    factory_return_types,
2450                                ) && engine.resolve_symbol(lhs).is_unknown()
2451                                {
2452                                    engine.seed_symbol(lhs.clone(), receiver_type.clone());
2453                                    changed = true;
2454                                    continue;
2455                                }
2456
2457                                if let Some(receiver_type) = normalized_receiver_type(callee)
2458                                    && engine.resolve_symbol(lhs).is_unknown()
2459                                {
2460                                    engine.seed_symbol(lhs.clone(), receiver_type);
2461                                    changed = true;
2462                                    continue;
2463                                }
2464                            }
2465                        }
2466                        AssignmentRhs::Symbol(rhs_symbol) => {
2467                            if !engine.is_shadowed(rhs_symbol)
2468                                && let Some(receiver_type) = normalized_receiver_type(rhs_symbol)
2469                                && engine.resolve_symbol(lhs).is_unknown()
2470                            {
2471                                engine.seed_symbol(lhs.clone(), receiver_type);
2472                                changed = true;
2473                                continue;
2474                            }
2475
2476                            if let SymbolResolution::Precise(targets) =
2477                                engine.resolve_symbol(rhs_symbol)
2478                                && !targets.is_empty()
2479                            {
2480                                aliases.push((lhs.clone(), rhs_symbol.clone()));
2481                            }
2482                        }
2483                        AssignmentRhs::Unknown => {}
2484                    }
2485                }
2486                ScopeFactEvent::Global { .. } | ScopeFactEvent::Nonlocal { .. } => {}
2487            }
2488        }
2489        let before = engine.snapshot();
2490        engine.apply_aliases_until_stable(aliases);
2491        if engine.snapshot() != before {
2492            changed = true;
2493        }
2494    }
2495
2496    engine.snapshot()
2497}
2498
2499fn factory_return_type_for_callee<'a>(
2500    callee: &str,
2501    current_class: Option<&str>,
2502    factory_return_types: &'a HashMap<String, String>,
2503) -> Option<&'a String> {
2504    if let Some(receiver_type) = factory_return_types.get(callee) {
2505        return Some(receiver_type);
2506    }
2507    let class_name = current_class?;
2508    let method = callee
2509        .strip_prefix("self.")
2510        .or_else(|| callee.strip_prefix("cls."))?;
2511    factory_return_types.get(&format!("{class_name}.{method}"))
2512}
2513
2514fn apply_annotation_event(
2515    symbol: &str,
2516    annotation: &str,
2517    allow_self_receivers: bool,
2518    engine: &mut LocalInferenceEngine<String>,
2519    changed: &mut bool,
2520) {
2521    if symbol.starts_with("self.") && !allow_self_receivers {
2522        return;
2523    }
2524    if let Some(receiver_type) = normalized_receiver_type(annotation)
2525        && engine.resolve_symbol(symbol).is_unknown()
2526    {
2527        engine.seed_symbol(symbol.to_string(), receiver_type);
2528        *changed = true;
2529    }
2530}
2531
2532enum ScopeFactEvent {
2533    Global {
2534        symbol: String,
2535    },
2536    Nonlocal {
2537        symbol: String,
2538    },
2539    Parameter {
2540        symbol: String,
2541        annotation: Option<String>,
2542    },
2543    Annotation {
2544        symbol: String,
2545        annotation: String,
2546    },
2547    Assignment {
2548        lhs: String,
2549        rhs: AssignmentRhs,
2550    },
2551}
2552
2553enum AssignmentRhs {
2554    Symbol(String),
2555    Call(String),
2556    Unknown,
2557}
2558
2559#[derive(Clone, Copy)]
2560enum ScopeFactTraversal {
2561    Module,
2562    Function,
2563    Class,
2564}
2565
2566fn collect_scope_fact_events(source: &str, traversal: ScopeFactTraversal) -> Vec<ScopeFactEvent> {
2567    if source.trim().is_empty() {
2568        return Vec::new();
2569    }
2570
2571    let mut parser = Parser::new();
2572    if parser
2573        .set_language(&tree_sitter_python::LANGUAGE.into())
2574        .is_err()
2575    {
2576        return Vec::new();
2577    }
2578    let Some(tree) = parser.parse(source, None) else {
2579        return Vec::new();
2580    };
2581
2582    let mut events = Vec::new();
2583    collect_scope_fact_events_from_node(tree.root_node(), source, traversal, &mut events);
2584    events
2585}
2586
2587fn collect_scope_fact_events_from_node(
2588    root: Node<'_>,
2589    source: &str,
2590    traversal: ScopeFactTraversal,
2591    events: &mut Vec<ScopeFactEvent>,
2592) {
2593    let mut stack = vec![(root, false)];
2594    while let Some((node, inside_function)) = stack.pop() {
2595        let next_inside_function = match traversal {
2596            ScopeFactTraversal::Module => {
2597                if matches!(
2598                    node.kind(),
2599                    "function_definition" | "class_definition" | "lambda"
2600                ) {
2601                    continue;
2602                }
2603                false
2604            }
2605            ScopeFactTraversal::Function => match node.kind() {
2606                "function_definition" if inside_function => continue,
2607                "function_definition" => true,
2608                "class_definition" | "lambda" => continue,
2609                _ => inside_function,
2610            },
2611            ScopeFactTraversal::Class => inside_function,
2612        };
2613        if matches!(traversal, ScopeFactTraversal::Function) && !next_inside_function {
2614            let mut cursor = node.walk();
2615            let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2616            children.reverse();
2617            stack.extend(children.into_iter().map(|child| (child, false)));
2618            continue;
2619        }
2620        match node.kind() {
2621            "global_statement" => collect_scope_directive_events(node, source, events, |symbol| {
2622                ScopeFactEvent::Global { symbol }
2623            }),
2624            "nonlocal_statement" => {
2625                collect_scope_directive_events(node, source, events, |symbol| {
2626                    ScopeFactEvent::Nonlocal { symbol }
2627                })
2628            }
2629            "parameters" | "lambda_parameters" => collect_parameter_events(node, source, events),
2630            "assignment" => collect_assignment_events(node, source, events),
2631            _ => {}
2632        }
2633
2634        let mut cursor = node.walk();
2635        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2636        children.reverse();
2637        stack.extend(
2638            children
2639                .into_iter()
2640                .map(|child| (child, next_inside_function)),
2641        );
2642    }
2643}
2644
2645fn collect_scope_directive_events(
2646    node: Node<'_>,
2647    source: &str,
2648    events: &mut Vec<ScopeFactEvent>,
2649    make_event: impl Fn(String) -> ScopeFactEvent,
2650) {
2651    let mut cursor = node.walk();
2652    for identifier in node
2653        .named_children(&mut cursor)
2654        .filter(|child| child.kind() == "identifier")
2655    {
2656        let Some(symbol) = non_empty_node_text(identifier, source) else {
2657            continue;
2658        };
2659        events.push(make_event(symbol));
2660    }
2661}
2662
2663fn collect_parameter_events(node: Node<'_>, source: &str, events: &mut Vec<ScopeFactEvent>) {
2664    let mut cursor = node.walk();
2665    for child in node.named_children(&mut cursor) {
2666        if child.kind() == "type_parameter" {
2667            continue;
2668        }
2669        let Some(symbol) = parameter_symbol(child, source) else {
2670            continue;
2671        };
2672        if matches!(symbol.as_str(), "self" | "cls" | "/") {
2673            continue;
2674        }
2675        let annotation = child
2676            .child_by_field_name("type")
2677            .map(|annotation| slice(annotation, source).trim().to_string())
2678            .filter(|annotation| !annotation.is_empty());
2679        events.push(ScopeFactEvent::Parameter { symbol, annotation });
2680    }
2681}
2682
2683fn parameter_symbol(node: Node<'_>, source: &str) -> Option<String> {
2684    if node.kind() == "identifier" {
2685        return non_empty_node_text(node, source);
2686    }
2687    if let Some(name) = node.child_by_field_name("name") {
2688        return non_empty_node_text(name, source);
2689    }
2690    let mut cursor = node.walk();
2691    node.named_children(&mut cursor)
2692        .find(|child| child.kind() == "identifier")
2693        .and_then(|identifier| non_empty_node_text(identifier, source))
2694}
2695
2696fn collect_assignment_events(node: Node<'_>, source: &str, events: &mut Vec<ScopeFactEvent>) {
2697    let Some(left) = node.child_by_field_name("left") else {
2698        return;
2699    };
2700    let Some(lhs) = receiver_symbol(left, source) else {
2701        return;
2702    };
2703
2704    if let Some(annotation) = node
2705        .child_by_field_name("type")
2706        .map(|annotation| slice(annotation, source).trim().to_string())
2707        .filter(|annotation| !annotation.is_empty())
2708    {
2709        events.push(ScopeFactEvent::Annotation {
2710            symbol: lhs,
2711            annotation,
2712        });
2713        return;
2714    }
2715
2716    let rhs = node
2717        .child_by_field_name("right")
2718        .and_then(|right| rhs_symbol(right, source))
2719        .unwrap_or(AssignmentRhs::Unknown);
2720    events.push(ScopeFactEvent::Assignment { lhs, rhs });
2721}
2722
2723fn receiver_symbol(node: Node<'_>, source: &str) -> Option<String> {
2724    match node.kind() {
2725        "identifier" | "attribute" => non_empty_node_text(node, source),
2726        _ => None,
2727    }
2728}
2729
2730fn rhs_symbol(node: Node<'_>, source: &str) -> Option<AssignmentRhs> {
2731    match node.kind() {
2732        "identifier" | "attribute" => non_empty_node_text(node, source).map(AssignmentRhs::Symbol),
2733        "call" => node
2734            .child_by_field_name("function")
2735            .or_else(|| node.named_child(0))
2736            .and_then(|callee| receiver_symbol(callee, source))
2737            .map(AssignmentRhs::Call),
2738        _ => None,
2739    }
2740}
2741
2742fn non_empty_node_text(node: Node<'_>, source: &str) -> Option<String> {
2743    let text = slice(node, source).trim();
2744    (!text.is_empty()).then(|| text.to_string())
2745}
2746
2747fn collect_factory_return_types_from_root(root: Node<'_>, source: &str) -> HashMap<String, String> {
2748    let mut returns = HashMap::default();
2749    let mut stack = vec![(root, None::<String>)];
2750    while let Some((node, class_name)) = stack.pop() {
2751        match node.kind() {
2752            "class_definition" => {
2753                let next_class = node
2754                    .child_by_field_name("name")
2755                    .and_then(|name| non_empty_node_text(name, source))
2756                    .or(class_name);
2757                push_factory_index_children(node, next_class, &mut stack);
2758            }
2759            "function_definition" => {
2760                if let Some(name) = node
2761                    .child_by_field_name("name")
2762                    .and_then(|name| non_empty_node_text(name, source))
2763                    && let Some(return_type) = factory_return_type(node, source)
2764                {
2765                    let key = class_name
2766                        .as_ref()
2767                        .map(|class| format!("{class}.{name}"))
2768                        .unwrap_or(name);
2769                    returns.insert(key, return_type);
2770                }
2771            }
2772            _ => push_factory_index_children(node, class_name, &mut stack),
2773        }
2774    }
2775    returns
2776}
2777
2778fn push_factory_index_children<'tree>(
2779    node: Node<'tree>,
2780    class_name: Option<String>,
2781    stack: &mut Vec<(Node<'tree>, Option<String>)>,
2782) {
2783    let mut cursor = node.walk();
2784    let mut children: Vec<Node<'tree>> = node.named_children(&mut cursor).collect();
2785    children.reverse();
2786    stack.extend(
2787        children
2788            .into_iter()
2789            .map(|child| (child, class_name.clone())),
2790    );
2791}
2792
2793fn factory_return_type(function: Node<'_>, source: &str) -> Option<String> {
2794    if let Some(return_type) = function.child_by_field_name("return_type") {
2795        return receiver_type_from_annotation_node(return_type, source);
2796    }
2797
2798    let body = function.child_by_field_name("body")?;
2799    let mut candidates = HashSet::default();
2800    let mut saw_return = false;
2801    let mut saw_unknown_return = false;
2802    let mut stack = vec![body];
2803    while let Some(node) = stack.pop() {
2804        if node != body && matches!(node.kind(), "function_definition" | "class_definition") {
2805            continue;
2806        }
2807        if node.kind() == "return_statement" {
2808            saw_return = true;
2809            match node
2810                .named_child(0)
2811                .and_then(|value| returned_receiver_type(value, source))
2812            {
2813                Some(returned_type) => {
2814                    candidates.insert(returned_type);
2815                }
2816                None => saw_unknown_return = true,
2817            }
2818        }
2819        let mut cursor = node.walk();
2820        let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2821        children.reverse();
2822        stack.extend(children);
2823    }
2824    if !saw_return || saw_unknown_return {
2825        return None;
2826    }
2827    (candidates.len() == 1)
2828        .then(|| candidates.into_iter().next())
2829        .flatten()
2830}
2831
2832/// Return the runtime class named by a structured Python return annotation.
2833///
2834/// For `Manager[A, B]`, the constructed class is the subscript base `Manager`.
2835/// `Optional[T]` is different: it denotes `T | None`, so retain the existing
2836/// supported-wrapper behavior and inspect its structured type argument.
2837fn receiver_type_from_annotation_node(annotation: Node<'_>, source: &str) -> Option<String> {
2838    match annotation.kind() {
2839        "type" => receiver_type_from_annotation_node(annotation.named_child(0)?, source),
2840        "identifier" | "attribute" | "member_type" | "string" => {
2841            normalized_receiver_type(slice(annotation, source).trim())
2842        }
2843        "generic_type" => {
2844            let base = annotation.named_child(0)?;
2845            if optional_annotation_wrapper(base, source) {
2846                let parameter = annotation.named_child(1)?;
2847                return receiver_type_from_annotation_node(parameter.named_child(0)?, source);
2848            }
2849            receiver_type_from_annotation_node(base, source)
2850        }
2851        "subscript" => {
2852            let value = annotation.child_by_field_name("value")?;
2853            if optional_annotation_wrapper(value, source) {
2854                let inner = annotation.child_by_field_name("subscript")?;
2855                return receiver_type_from_annotation_node(inner, source);
2856            }
2857            normalized_receiver_type(slice(value, source).trim())
2858        }
2859        _ => None,
2860    }
2861}
2862
2863fn optional_annotation_wrapper(node: Node<'_>, source: &str) -> bool {
2864    match node.kind() {
2865        "identifier" => slice(node, source) == "Optional",
2866        "attribute" => {
2867            let (Some(object), Some(attribute)) = (
2868                node.child_by_field_name("object"),
2869                node.child_by_field_name("attribute"),
2870            ) else {
2871                return false;
2872            };
2873            object.kind() == "identifier"
2874                && attribute.kind() == "identifier"
2875                && slice(object, source) == "typing"
2876                && slice(attribute, source) == "Optional"
2877        }
2878        _ => false,
2879    }
2880}
2881
2882fn returned_receiver_type(node: Node<'_>, source: &str) -> Option<String> {
2883    let raw = match node.kind() {
2884        "identifier" => non_empty_node_text(node, source),
2885        "call" => node
2886            .child_by_field_name("function")
2887            .or_else(|| node.named_child(0))
2888            .filter(|callee| callee.kind() == "identifier")
2889            .and_then(|callee| non_empty_node_text(callee, source)),
2890        _ => None,
2891    }?;
2892    normalized_receiver_type(&raw)
2893}
2894
2895#[cfg(test)]
2896mod tests {
2897    use super::*;
2898    use std::path::PathBuf;
2899
2900    #[test]
2901    fn lambda_scope_facts_preserve_untyped_parameter_shadowing() {
2902        let source = "shadowed = lambda method: method.signature\n";
2903        let mut parser = Parser::new();
2904        parser
2905            .set_language(&tree_sitter_python::LANGUAGE.into())
2906            .unwrap();
2907        let tree = parser.parse(source, None).unwrap();
2908        let mut nodes = vec![tree.root_node()];
2909        let lambda = loop {
2910            let node = nodes.pop().unwrap();
2911            if node.kind() == "lambda" {
2912                break node;
2913            }
2914            let mut cursor = node.walk();
2915            nodes.extend(node.named_children(&mut cursor));
2916        };
2917
2918        let facts = collect_function_scope_facts_from_node(lambda, source);
2919
2920        assert!(facts.is_shadowed("method"));
2921        assert!(facts.resolution_for("method").is_unknown());
2922    }
2923
2924    #[test]
2925    fn pre_cancelled_graph_build_skips_python_file_parsing() {
2926        let temp = tempfile::tempdir().unwrap();
2927        let root = temp.path().canonicalize().unwrap();
2928        std::fs::write(root.join("target.py"), "def target():\n    pass\n").unwrap();
2929        let file = ProjectFile::new(root.clone(), PathBuf::from("target.py"));
2930        let files = [file.clone()].into_iter().collect();
2931        let cancellation = CancellationToken::default();
2932        cancellation.cancel();
2933
2934        let graph = build_python_graph(&files, &file, Some(&cancellation));
2935
2936        assert!(graph.parsed.is_empty());
2937    }
2938
2939    #[test]
2940    fn graph_build_parses_only_candidates_and_target_not_transitive_imports() {
2941        let temp = tempfile::tempdir().unwrap();
2942        let root = temp.path().canonicalize().unwrap();
2943        std::fs::write(root.join("target.py"), "from dependency import value\n").unwrap();
2944        std::fs::write(
2945            root.join("candidate.py"),
2946            "from transitively_imported import value\n",
2947        )
2948        .unwrap();
2949        std::fs::write(root.join("dependency.py"), "value = 1\n").unwrap();
2950        std::fs::write(root.join("transitively_imported.py"), "value = 2\n").unwrap();
2951        let target = ProjectFile::new(root.clone(), PathBuf::from("target.py"));
2952        let candidate = ProjectFile::new(root.clone(), PathBuf::from("candidate.py"));
2953        let dependency = ProjectFile::new(root.clone(), PathBuf::from("dependency.py"));
2954        let transitive = ProjectFile::new(root.clone(), PathBuf::from("transitively_imported.py"));
2955        let candidates = [candidate.clone()].into_iter().collect();
2956
2957        let graph = build_python_graph(&candidates, &target, None);
2958
2959        assert_eq!(graph.parsed.len(), 2);
2960        assert!(graph.parsed.contains_key(&target));
2961        assert!(graph.parsed.contains_key(&candidate));
2962        assert!(!graph.parsed.contains_key(&dependency));
2963        assert!(!graph.parsed.contains_key(&transitive));
2964    }
2965}