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