1use brokk_bifrost_core::analyzer::capabilities::ImportAnalysisProvider;
13use brokk_bifrost_core::analyzer::model::ImportInfo;
14use brokk_bifrost_core::analyzer::prepared_syntax::{IndexedFileFacts, PreparedSyntaxTree};
15use brokk_bifrost_core::analyzer::project::Project;
16use brokk_bifrost_core::analyzer::query_token::QueryToken;
17use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
18use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
19use brokk_bifrost_core::analyzer::usages::model::{
20 ExportEntry, ExportIndex, ImportBinder, ImportBinding, ImportKind, ReexportStar,
21};
22use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, Language, ProjectFile};
23use brokk_bifrost_core::hash::{HashMap, HashSet};
24use std::collections::BTreeSet;
25use std::path::{Path, PathBuf};
26use std::sync::Arc;
27use tree_sitter::Node;
28
29use crate::bindings::{python_direct_scope_bindings_bounded, python_executable_scope_children};
30use crate::declarations::{
31 collect_python_identifiers, module_code_unit, parse_python_tree, python_base_origin_node,
32 python_declared_import_root, python_import_root, python_project_module_components,
33 python_repeated_root_extension_dir,
34};
35use crate::imports::{
36 PythonImportDetails, python_import_details, python_namespace_binding_module,
37 python_namespace_binding_name, resolve_exported_fqn, resolve_import_bindings,
38 resolve_python_relative_module,
39};
40use crate::syntax::{python_plain_string_literal, python_static_attribute_path};
41use crate::usage_index::PythonUsageIndex;
42
43pub trait PythonSource:
54 CodeUnitIndex + ImportAnalysisProvider + crate::source_facts::PythonSourceFactProvider
55{
56 fn path_module_fqn(&self, module_fq: &str) -> Option<Vec<CodeUnit>>;
59
60 fn path_module_fqns_batch(&self, module_fqs: &[String]) -> Vec<Option<Vec<CodeUnit>>>;
63
64 fn definition_fqn(&self, fqn: &str) -> Vec<CodeUnit>;
65
66 fn module_spellings(&self) -> Arc<PythonModuleSpellings>;
73
74 fn import_binder_of(&self, file: &ProjectFile) -> Arc<ImportBinder>;
79
80 fn export_index_of(&self, file: &ProjectFile) -> Arc<ExportIndex>;
81
82 fn prepared_syntax(
87 &self,
88 token: QueryToken<'_>,
89 file: &ProjectFile,
90 ) -> Option<Arc<PreparedSyntaxTree>>;
91
92 fn visit_file_facts(
95 &self,
96 files: &[ProjectFile],
97 visit: &mut dyn FnMut(&ProjectFile, Option<&dyn IndexedFileFacts>),
98 );
99}
100
101pub trait PythonUsageSource: PythonSource {
104 fn usage_index(&self) -> Arc<PythonUsageIndex>;
105}
106
107pub fn extract_type_identifiers(source: &str) -> BTreeSet<String> {
108 let Some(tree) = parse_python_tree(source) else {
109 return BTreeSet::new();
110 };
111 let mut identifiers = HashSet::default();
112 collect_python_identifiers(tree.root_node(), source, &mut identifiers);
113 identifiers.into_iter().collect()
114}
115
116pub fn retain_modules_for_importer<T>(
135 project: &dyn Project,
136 importer: &ProjectFile,
137 candidates: &mut Vec<T>,
138 source_of: impl Fn(&T) -> &ProjectFile,
139) {
140 retain_nearest_owning_root(importer, candidates, |candidate| {
141 let source = source_of(candidate);
142 debug_assert_eq!(
143 importer.root(),
144 source.root(),
145 "one workspace's Python files share one project root"
146 );
147 python_import_root(project, source)
148 });
149}
150
151fn retain_nearest_owning_root<T>(
160 importer: &ProjectFile,
161 candidates: &mut Vec<T>,
162 root_of: impl Fn(&T) -> PathBuf,
163) {
164 if candidates.len() < 2 {
165 return;
166 }
167 let owning_depth = |candidate: &T| {
168 let root = root_of(candidate);
169 importer
170 .rel_path()
171 .starts_with(&root)
172 .then(|| root.components().count())
173 };
174 let depths = candidates.iter().map(owning_depth).collect::<Vec<_>>();
175 let Some(nearest) = depths.iter().copied().flatten().max() else {
176 return;
177 };
178 *candidates = std::mem::take(candidates)
179 .into_iter()
180 .zip(depths)
181 .filter(|(_, depth)| *depth == Some(nearest))
182 .map(|(candidate, _)| candidate)
183 .collect();
184}
185
186#[derive(Clone, Debug)]
188struct ProjectRootModuleName {
189 file: ProjectFile,
190 components: Vec<String>,
191}
192
193#[derive(Clone, Debug)]
202pub struct PythonSourceRootModule {
203 pub file: ProjectFile,
204 pub source_root: PathBuf,
208}
209
210#[derive(Default)]
232pub struct PythonModuleSpellings {
233 project: Option<Arc<dyn Project>>,
234 by_last_component: HashMap<String, Vec<ProjectRootModuleName>>,
235 package_dirs: HashSet<PathBuf>,
238 path_extensions: HashMap<PathBuf, HashSet<PathBuf>>,
250}
251
252impl std::fmt::Debug for PythonModuleSpellings {
253 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254 f.debug_struct("PythonModuleSpellings")
255 .field("by_last_component", &self.by_last_component)
256 .field("package_dirs", &self.package_dirs)
257 .field("path_extensions", &self.path_extensions)
258 .finish_non_exhaustive()
259 }
260}
261
262impl PythonModuleSpellings {
263 pub(crate) fn project(&self) -> Option<&dyn Project> {
264 self.project.as_deref()
265 }
266
267 pub fn build(files: impl IntoIterator<Item = ProjectFile>, project: Arc<dyn Project>) -> Self {
273 let mut by_last_component: HashMap<String, Vec<ProjectRootModuleName>> = HashMap::default();
274 let mut package_dirs: HashSet<PathBuf> = HashSet::default();
275 let mut package_markers: Vec<ProjectFile> = Vec::new();
276 for file in files {
277 if file.rel_path().file_name().and_then(|name| name.to_str()) == Some("__init__.py")
278 && let Some(directory) = file.rel_path().parent()
279 {
280 package_dirs.insert(directory.to_path_buf());
281 package_markers.push(file.clone());
282 }
283 let components = python_project_module_components(&file);
284 let Some(last) = components.last() else {
285 continue;
286 };
287 by_last_component
288 .entry(last.clone())
289 .or_default()
290 .push(ProjectRootModuleName { file, components });
291 }
292 let mut path_extensions: HashMap<PathBuf, HashSet<PathBuf>> = HashMap::default();
293 for marker in package_markers {
294 let Some(directory) = marker.rel_path().parent() else {
295 continue;
296 };
297 if !python_package_holds_a_repeat_of_itself(directory, &package_dirs) {
298 continue;
299 }
300 let Ok(source) = project.read_source(&marker) else {
301 continue;
302 };
303 let extended = python_appended_path_dirs(&source, marker.rel_path());
304 if !extended.is_empty() {
305 path_extensions
306 .entry(directory.to_path_buf())
307 .or_default()
308 .extend(extended);
309 }
310 }
311 Self {
312 project: Some(project),
313 by_last_component,
314 package_dirs,
315 path_extensions,
316 }
317 }
318
319 pub fn suffix_candidates(&self, module_fq: &str) -> Vec<PythonSourceRootModule> {
331 if module_fq.starts_with('.') {
335 return Vec::new();
336 }
337 let components = parse_symbol_path(Language::Python, module_fq);
338 let Some(last) = components.last().filter(|_| components.len() > 1) else {
339 return Vec::new();
340 };
341 let mut modules = Vec::new();
342 for candidate in self
343 .by_last_component
344 .get(last.as_str())
345 .into_iter()
346 .flatten()
347 {
348 let Some(dropped) = candidate
351 .components
352 .len()
353 .checked_sub(components.len())
354 .filter(|dropped| *dropped > 0)
355 else {
356 continue;
357 };
358 if candidate.components[dropped..] != components[..] {
359 continue;
360 }
361 let source_root: PathBuf = candidate
364 .file
365 .rel_path()
366 .components()
367 .take(dropped)
368 .collect();
369 debug_assert!(
370 !source_root.as_os_str().is_empty(),
371 "dropping a component of {:?} leaves a directory",
372 candidate.file.rel_path()
373 );
374 modules.push(PythonSourceRootModule {
375 file: candidate.file.clone(),
376 source_root,
377 });
378 }
379 modules
380 }
381
382 pub fn establishes_source_root(&self, candidate: &PythonSourceRootModule) -> bool {
428 if let Some(declared) = python_declared_import_root(
429 self.project()
430 .expect("source-root candidates have a selected project"),
431 &candidate.file,
432 ) {
433 if declared == candidate.source_root {
434 return true;
435 }
436 let Some(extension) = python_repeated_root_extension_dir(
437 &candidate.file,
438 &candidate.source_root,
439 &declared,
440 ) else {
441 return false;
442 };
443 return self.package_chain_reaches_module(candidate)
444 && self
445 .path_extensions
446 .get(&candidate.source_root)
447 .is_some_and(|extended| extended.contains(&extension));
448 }
449 self.package_chain_reaches_module(candidate)
450 }
451
452 fn package_chain_reaches_module(&self, candidate: &PythonSourceRootModule) -> bool {
461 debug_assert!(
462 candidate
463 .file
464 .rel_path()
465 .starts_with(&candidate.source_root),
466 "a candidate's implied root is a prefix of its own path"
467 );
468 let mut directory = candidate.file.rel_path().parent();
469 while let Some(current) = directory {
470 if !self.package_dirs.contains(current) {
471 return false;
472 }
473 if current == candidate.source_root {
474 return true;
475 }
476 directory = current.parent();
477 }
478 false
479 }
480}
481
482fn python_package_holds_a_repeat_of_itself(
490 directory: &Path,
491 package_dirs: &HashSet<PathBuf>,
492) -> bool {
493 directory
494 .file_name()
495 .is_some_and(|name| package_dirs.contains(&directory.join(name)))
496}
497
498fn python_appended_path_dirs(source: &str, marker_rel_path: &Path) -> Vec<PathBuf> {
513 let Some(tree) = parse_python_tree(source) else {
514 return Vec::new();
515 };
516 let mut bound: HashMap<String, PathBuf> = HashMap::default();
517 bound.insert("__file__".to_owned(), marker_rel_path.to_path_buf());
518 let mut appended: Vec<PathBuf> = Vec::new();
519 walk_named_tree_preorder(tree.root_node(), true, |node| {
520 match node.kind() {
521 "function_definition" | "class_definition" | "lambda" => {
525 return WalkControl::SkipChildren;
526 }
527 "assignment" => {
528 if python_names_path(node.child_by_field_name("left"), source) {
533 appended.clear();
534 }
535 python_bind_assignment(node, source, &mut bound);
536 }
537 "call" => {
538 let directories = python_path_appended_dirs_in_call(node, source, &bound);
539 if directories.is_empty()
540 || !python_extension_statement_runs(node, source, &bound, &directories)
541 {
542 return WalkControl::Continue;
543 }
544 for directory in directories {
545 if directory.is_relative()
546 && !directory.as_os_str().is_empty()
547 && !appended.contains(&directory)
548 {
549 appended.push(directory);
550 }
551 }
552 }
553 _ => {}
554 }
555 WalkControl::Continue
556 });
557 appended
558}
559
560fn python_names_path(node: Option<Node<'_>>, source: &str) -> bool {
562 node.is_some_and(|node| {
563 node.kind() == "identifier" && node.utf8_text(source.as_bytes()).ok() == Some("__path__")
564 })
565}
566
567fn python_extension_statement_runs(
582 call: Node<'_>,
583 source: &str,
584 bound: &HashMap<String, PathBuf>,
585 directories: &[PathBuf],
586) -> bool {
587 let Some(statement) = call.parent() else {
588 return false;
589 };
590 if statement.kind() != "expression_statement" {
591 return false;
592 }
593 let Some(container) = statement.parent() else {
594 return false;
595 };
596 if container.kind() == "module" {
597 return true;
598 }
599 let guarded = if container.kind() == "block" {
604 container.parent().filter(|parent| {
605 parent.kind() == "if_statement"
606 && parent.child_by_field_name("consequence") == Some(container)
607 })
608 } else if container.kind() == "if_statement"
609 && container.child_by_field_name("consequence") == Some(statement)
610 {
611 Some(container)
612 } else {
613 None
614 };
615 let Some(guarded) = guarded else {
616 return false;
617 };
618 if guarded.parent().map(|parent| parent.kind()) != Some("module") {
623 return false;
624 }
625 let Some(condition) = guarded.child_by_field_name("condition") else {
626 return false;
627 };
628 python_condition_justifies_path_extension(condition, source, bound, directories)
629}
630
631fn python_condition_justifies_path_extension(
636 condition: Node<'_>,
637 source: &str,
638 bound: &HashMap<String, PathBuf>,
639 directories: &[PathBuf],
640) -> bool {
641 if condition.kind() != "comparison_operator" {
642 return false;
643 }
644 let mut cursor = condition.walk();
645 let operators = condition
646 .children_by_field_name("operators", &mut cursor)
647 .map(|operator| operator.kind())
648 .collect::<Vec<_>>();
649 if operators.as_slice() != ["not in"] {
650 return false;
651 }
652 let mut cursor = condition.walk();
655 let operands: Vec<Node<'_>> = condition.named_children(&mut cursor).collect();
656 let [left, right] = operands.as_slice() else {
657 return false;
658 };
659 if !python_names_path(Some(*right), source) {
660 return false;
661 }
662 python_path_expression_dir(*left, source, bound)
663 .is_some_and(|directory| directories.contains(&directory))
664}
665
666fn python_bind_assignment(
671 assignment: Node<'_>,
672 source: &str,
673 bound: &mut HashMap<String, PathBuf>,
674) {
675 let mut targets: Vec<Node<'_>> = Vec::new();
676 let mut current = assignment;
677 loop {
678 let (Some(left), Some(right)) = (
679 current.child_by_field_name("left"),
680 current.child_by_field_name("right"),
681 ) else {
682 return;
683 };
684 targets.push(left);
685 if right.kind() != "assignment" {
686 let resolved = python_path_expression_dir(right, source, bound);
687 for target in targets {
688 if target.kind() != "identifier" {
689 continue;
690 }
691 let Ok(name) = target.utf8_text(source.as_bytes()) else {
692 continue;
693 };
694 match &resolved {
695 Some(directory) => {
696 bound.insert(name.to_owned(), directory.clone());
697 }
698 None => {
699 bound.remove(name);
700 }
701 }
702 }
703 return;
704 }
705 current = right;
706 }
707}
708
709fn python_path_appended_dirs_in_call(
716 call: Node<'_>,
717 source: &str,
718 bound: &HashMap<String, PathBuf>,
719) -> Vec<PathBuf> {
720 let Some(function) = call.child_by_field_name("function") else {
721 return Vec::new();
722 };
723 let Some(path) = python_static_attribute_path(function) else {
724 return Vec::new();
725 };
726 let [object, method] = path.as_slice() else {
727 return Vec::new();
728 };
729 if object.utf8_text(source.as_bytes()).ok() != Some("__path__") {
730 return Vec::new();
731 }
732 let Ok(method) = method.utf8_text(source.as_bytes()) else {
733 return Vec::new();
734 };
735 let Some(arguments) = call.child_by_field_name("arguments") else {
736 return Vec::new();
737 };
738 let mut cursor = arguments.walk();
739 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
740 let named: Vec<Node<'_>> = match method {
741 "append" => arguments.first().copied().into_iter().collect(),
742 "insert" => arguments.get(1).copied().into_iter().collect(),
743 "extend" => match arguments.first() {
744 Some(list) if list.kind() == "list" => {
745 let mut cursor = list.walk();
746 list.named_children(&mut cursor).collect()
747 }
748 Some(iterator) => vec![*iterator],
749 None => Vec::new(),
750 },
751 _ => Vec::new(),
752 };
753 named
754 .into_iter()
755 .filter_map(|node| python_path_expression_dir(node, source, bound))
756 .collect()
757}
758
759fn python_path_expression_dir(
763 node: Node<'_>,
764 source: &str,
765 bound: &HashMap<String, PathBuf>,
766) -> Option<PathBuf> {
767 if let Some(literal) = python_plain_string_literal(node, source) {
768 let directory = PathBuf::from(literal);
769 return (!directory.as_os_str().is_empty()).then_some(directory);
770 }
771 match node.kind() {
772 "identifier" => bound.get(node.utf8_text(source.as_bytes()).ok()?).cloned(),
773 "call" => python_path_call_dir(node, source, bound),
774 _ => None,
775 }
776}
777
778fn python_path_call_dir(
781 call: Node<'_>,
782 source: &str,
783 bound: &HashMap<String, PathBuf>,
784) -> Option<PathBuf> {
785 let path = python_static_attribute_path(call.child_by_field_name("function")?)?;
786 let names: Vec<&str> = path
787 .iter()
788 .map(|node| node.utf8_text(source.as_bytes()).ok())
789 .collect::<Option<Vec<_>>>()?;
790 let arguments = call.child_by_field_name("arguments")?;
791 let mut cursor = arguments.walk();
792 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
793 match names.as_slice() {
794 ["os", "path", "dirname"] => {
795 let directory = python_path_expression_dir(*arguments.first()?, source, bound)?;
796 Some(directory.parent()?.to_path_buf())
797 }
798 ["os", "path", "join"] => {
799 let mut parts = arguments.iter();
800 let mut joined = python_path_expression_dir(*parts.next()?, source, bound)?;
801 for part in parts {
802 joined.push(python_path_expression_dir(*part, source, bound)?);
803 }
804 Some(joined)
805 }
806 _ => None,
807 }
808}
809
810pub fn resolve_module_code_unit(
818 python: &dyn PythonSource,
819 importer: Option<&ProjectFile>,
820 module_fq: &str,
821) -> Option<CodeUnit> {
822 let indexed = match python.path_module_fqn(module_fq) {
823 Some(units) => indexed_module_for_importer(python.project(), importer, units),
826 None => indexed_module_for_importer(
827 python.project(),
828 importer,
829 python.definition_fqn(module_fq),
830 ),
831 };
832 match indexed {
833 IndexedModuleLookup::Unique(module) => Some(module),
834 IndexedModuleLookup::Ambiguous => None,
838 IndexedModuleLookup::Missing => {
839 let importer = importer?;
840 module_under_source_root(&python.module_spellings(), importer, module_fq)
841 }
842 }
843}
844
845enum IndexedModuleLookup {
854 Unique(CodeUnit),
855 Ambiguous,
856 Missing,
857}
858
859fn indexed_module_for_importer(
860 project: &dyn Project,
861 importer: Option<&ProjectFile>,
862 units: Vec<CodeUnit>,
863) -> IndexedModuleLookup {
864 let mut modules = units
865 .into_iter()
866 .filter(CodeUnit::is_module)
867 .collect::<Vec<_>>();
868 if let Some(importer) = importer {
869 retain_modules_for_importer(project, importer, &mut modules, CodeUnit::source);
870 }
871 match modules.len() {
872 0 => IndexedModuleLookup::Missing,
873 1 => IndexedModuleLookup::Unique(modules.remove(0)),
874 _ => IndexedModuleLookup::Ambiguous,
875 }
876}
877
878pub fn python_module_file_under_source_root(
907 spellings: &PythonModuleSpellings,
908 importer: &ProjectFile,
909 module_fq: &str,
910) -> Option<ProjectFile> {
911 let mut candidates = spellings.suffix_candidates(module_fq);
912 retain_nearest_owning_root(importer, &mut candidates, |candidate| {
913 candidate.source_root.clone()
914 });
915 candidates.retain(|candidate| spellings.establishes_source_root(candidate));
916 let [candidate] = candidates.as_slice() else {
917 return None;
918 };
919 debug_assert_eq!(
920 importer.root(),
921 candidate.file.root(),
922 "one workspace's Python files share one project root"
923 );
924 Some(candidate.file.clone())
925}
926
927fn module_under_source_root(
932 spellings: &PythonModuleSpellings,
933 importer: &ProjectFile,
934 module_fq: &str,
935) -> Option<CodeUnit> {
936 let file = python_module_file_under_source_root(spellings, importer, module_fq)?;
937 module_code_unit(spellings.project()?, &file)
938}
939
940pub fn resolve_module_code_units_batch(
946 python: &dyn PythonSource,
947 importer: Option<&ProjectFile>,
948 module_fqs: &[String],
949) -> Vec<Option<CodeUnit>> {
950 let path_results = python.path_module_fqns_batch(module_fqs);
951 let mut lookups: Vec<IndexedModuleLookup> = Vec::with_capacity(module_fqs.len());
952 let mut needs_definition_fallback = Vec::new();
953 for (i, units) in path_results.into_iter().enumerate() {
954 match units {
955 Some(units) => lookups.push(indexed_module_for_importer(
956 python.project(),
957 importer,
958 units,
959 )),
960 None => {
961 needs_definition_fallback.push(i);
962 lookups.push(IndexedModuleLookup::Missing);
964 }
965 }
966 }
967 for i in needs_definition_fallback {
968 lookups[i] = indexed_module_for_importer(
969 python.project(),
970 importer,
971 python.definition_fqn(&module_fqs[i]),
972 );
973 }
974 let mut results: Vec<Option<CodeUnit>> = vec![None; module_fqs.len()];
975 let mut needs_source_root = Vec::new();
976 for (i, lookup) in lookups.into_iter().enumerate() {
977 match lookup {
978 IndexedModuleLookup::Unique(module) => results[i] = Some(module),
979 IndexedModuleLookup::Ambiguous => {}
980 IndexedModuleLookup::Missing => needs_source_root.push(i),
981 }
982 }
983 if let Some(importer) = importer
984 && !needs_source_root.is_empty()
985 {
986 let spellings = python.module_spellings();
989 for i in needs_source_root {
990 results[i] = module_under_source_root(&spellings, importer, &module_fqs[i]);
991 }
992 }
993 results
994}
995
996pub fn compute_export_index_of(
997 python: &dyn PythonSource,
998 token: QueryToken<'_>,
999 file: &ProjectFile,
1000) -> ExportIndex {
1001 let mut index = ExportIndex::empty();
1002 let mut events = Vec::new();
1003 let declarations = python.top_level_declarations(file);
1004 collect_local_export_events(
1005 declarations.iter(),
1006 |code_unit| {
1007 python
1008 .ranges(code_unit)
1009 .iter()
1010 .map(|range| range.start_byte)
1011 .min()
1012 .unwrap_or(usize::MAX)
1013 },
1014 &mut events,
1015 );
1016
1017 let imports = python.import_info_of(token, file);
1018 collect_reexport_events_from_imports(python, file, &imports, &mut events, &mut index);
1019
1020 finish_export_index(events, index)
1021}
1022
1023pub fn export_index_from_file_facts(
1024 python: &dyn PythonSource,
1025 file: &ProjectFile,
1026 facts: &dyn IndexedFileFacts,
1027 module_name: &str,
1028) -> ExportIndex {
1029 let mut index = ExportIndex::empty();
1030 let mut events = Vec::new();
1031 collect_local_export_events(
1032 facts.top_level_declarations().iter(),
1033 |code_unit| {
1034 facts
1035 .declaration_ranges(code_unit)
1036 .into_iter()
1037 .flatten()
1038 .map(|range| range.start_byte)
1039 .min()
1040 .unwrap_or(usize::MAX)
1041 },
1042 &mut events,
1043 );
1044
1045 if !facts
1046 .top_level_declarations()
1047 .iter()
1048 .any(CodeUnit::is_module)
1049 && let Some(identifier) = module_name.rsplit('.').next()
1050 && !identifier.is_empty()
1051 && !identifier.starts_with('_')
1052 {
1053 events.push((
1054 0,
1055 identifier.to_string(),
1056 ExportEntry::Local {
1057 local_name: identifier.to_string(),
1058 },
1059 ));
1060 }
1061
1062 collect_reexport_events_from_imports(python, file, facts.imports(), &mut events, &mut index);
1063
1064 finish_export_index(events, index)
1065}
1066
1067fn collect_local_export_events<'a>(
1068 declarations: impl IntoIterator<Item = &'a CodeUnit>,
1069 mut start_byte: impl FnMut(&CodeUnit) -> usize,
1070 events: &mut Vec<(usize, String, ExportEntry)>,
1071) -> HashSet<String> {
1072 let mut local_names = HashSet::default();
1073 for code_unit in declarations {
1074 let identifier = code_unit.identifier().trim();
1075 if identifier.is_empty() {
1076 continue;
1077 }
1078 local_names.insert(identifier.to_string());
1079 events.push((
1080 start_byte(code_unit),
1081 identifier.to_string(),
1082 ExportEntry::Local {
1083 local_name: identifier.to_string(),
1084 },
1085 ));
1086 }
1087 local_names
1088}
1089
1090fn finish_export_index(
1091 mut events: Vec<(usize, String, ExportEntry)>,
1092 mut index: ExportIndex,
1093) -> ExportIndex {
1094 events.sort_by_key(|(start_byte, _, _)| *start_byte);
1095 for (_, exported_name, entry) in events {
1096 index.exports_by_name.insert(exported_name, entry);
1097 }
1098 index
1099}
1100
1101fn collect_reexport_events_from_imports(
1102 python: &dyn PythonSource,
1103 file: &ProjectFile,
1104 imports: &[ImportInfo],
1105 events: &mut Vec<(usize, String, ExportEntry)>,
1106 index: &mut ExportIndex,
1107) {
1108 for import in imports {
1109 if import
1110 .path
1111 .as_ref()
1112 .is_none_or(|path| !path.lexical_scopes.is_empty())
1113 {
1114 continue;
1115 }
1116 record_single_reexport_event(python, file, import, events, index);
1117 }
1118}
1119
1120fn record_single_reexport_event(
1121 python: &dyn PythonSource,
1122 file: &ProjectFile,
1123 import: &ImportInfo,
1124 events: &mut Vec<(usize, String, ExportEntry)>,
1125 index: &mut ExportIndex,
1126) {
1127 let Some(PythonImportDetails::FromImport {
1128 module,
1129 name,
1130 alias,
1131 wildcard,
1132 }) = python_import_details(import)
1133 else {
1134 return;
1135 };
1136 let start_byte = import
1137 .path
1138 .as_ref()
1139 .map(|path| path.declaration_start_byte)
1140 .unwrap_or(usize::MAX);
1141 let resolved_module = if module.starts_with('.') {
1142 resolve_python_relative_module(python.project(), file, &module)
1143 } else {
1144 Some(module.clone())
1145 };
1146 let Some(resolved_module) = resolved_module else {
1147 return;
1148 };
1149
1150 if wildcard {
1151 index.reexport_stars.push(ReexportStar {
1152 module_specifier: resolved_module,
1153 });
1154 return;
1155 }
1156 let exported_name = alias.unwrap_or(name.clone());
1157 let module_candidate = format!("{resolved_module}.{name}");
1163 if resolve_module_code_unit(python, Some(file), &module_candidate).is_some() {
1164 events.push((
1165 start_byte,
1166 exported_name,
1167 ExportEntry::ReexportedModule {
1168 module_specifier: module_candidate,
1169 },
1170 ));
1171 return;
1172 }
1173 events.push((
1174 start_byte,
1175 exported_name,
1176 ExportEntry::ReexportedNamed {
1177 module_specifier: resolved_module,
1178 imported_name: name,
1179 },
1180 ));
1181}
1182
1183pub fn import_binder_from_imports(
1184 python: &dyn PythonSource,
1185 file: &ProjectFile,
1186 imports: &[ImportInfo],
1187) -> ImportBinder {
1188 let mut binder = ImportBinder::empty();
1189
1190 for (local_name, binding) in import_bindings_from_imports(python, file, imports) {
1191 binder.bindings.insert(local_name, binding);
1192 }
1193
1194 binder
1195}
1196
1197pub fn import_bindings_from_imports(
1202 python: &dyn PythonSource,
1203 file: &ProjectFile,
1204 imports: &[ImportInfo],
1205) -> Vec<(String, ImportBinding)> {
1206 let mut bindings = Vec::new();
1207
1208 for import in imports {
1209 let Some(details) = python_import_details(import) else {
1210 continue;
1211 };
1212 match details {
1213 PythonImportDetails::Import { module, alias } => {
1214 let local_name = python_namespace_binding_name(import, alias.as_deref(), &module);
1215 let module_specifier =
1216 python_namespace_binding_module(import, alias.as_deref(), &module);
1217 bindings.push((
1218 local_name,
1219 ImportBinding {
1220 module_specifier,
1221 namespace_imported_module: Some(module),
1222 kind: ImportKind::Namespace,
1223 imported_name: None,
1224 },
1225 ));
1226 }
1227 PythonImportDetails::FromImport {
1228 module,
1229 name,
1230 wildcard,
1231 ..
1232 } => {
1233 let resolved_module = if module.starts_with('.') {
1234 resolve_python_relative_module(python.project(), file, &module)
1235 } else {
1236 Some(module.clone())
1237 };
1238 let Some(resolved_module) = resolved_module else {
1239 continue;
1240 };
1241 if wildcard {
1242 bindings.extend(
1247 public_declarations_in_module(python, file, &resolved_module)
1248 .into_iter()
1249 .map(|declaration| {
1250 let name = declaration.identifier().to_string();
1251 (
1252 name.clone(),
1253 ImportBinding {
1254 module_specifier: resolved_module.clone(),
1255 namespace_imported_module: None,
1256 kind: ImportKind::Named,
1257 imported_name: Some(name),
1258 },
1259 )
1260 }),
1261 );
1262 continue;
1263 }
1264 let local_name = import
1269 .local_name()
1270 .map(str::to_string)
1271 .unwrap_or_else(|| name.clone());
1272 let module_candidate = format!("{resolved_module}.{name}");
1273 let exported = resolve_exported_fqn(python, Some(file), &module_candidate);
1274 let exported_module =
1275 if !exported.is_empty() && exported.iter().all(CodeUnit::is_module) {
1276 let mut paths = exported.iter().map(CodeUnit::fq_segment_texts);
1277 let path = paths
1278 .next()
1279 .expect("non-empty exported modules have a first path");
1280 paths
1281 .all(|candidate| candidate == path)
1282 .then(|| path.join("."))
1283 } else {
1284 None
1285 };
1286 let implicit_submodule = if exported.is_empty()
1287 && resolve_module_code_unit(python, Some(file), &module_candidate).is_some()
1288 {
1289 Some(module_candidate)
1290 } else {
1291 None
1292 };
1293 if let Some(module_specifier) = exported_module.or(implicit_submodule) {
1294 bindings.push((
1295 local_name,
1296 ImportBinding {
1297 module_specifier,
1298 namespace_imported_module: None,
1299 kind: ImportKind::Namespace,
1300 imported_name: None,
1301 },
1302 ));
1303 continue;
1304 }
1305 bindings.push((
1306 local_name,
1307 ImportBinding {
1308 module_specifier: resolved_module,
1309 namespace_imported_module: None,
1310 kind: ImportKind::Named,
1311 imported_name: Some(name),
1312 },
1313 ));
1314 }
1315 }
1316 }
1317
1318 bindings
1319}
1320
1321pub fn public_declarations_in_module(
1322 python: &dyn PythonSource,
1323 importer: &ProjectFile,
1324 module_fq: &str,
1325) -> Vec<CodeUnit> {
1326 let Some(module_code_unit) = resolve_module_code_unit(python, Some(importer), module_fq) else {
1327 return Vec::new();
1328 };
1329 python
1330 .direct_children(&module_code_unit)
1331 .into_iter()
1332 .filter(|code_unit| !code_unit.identifier().starts_with('_'))
1333 .collect()
1334}
1335
1336pub fn resolve_base_class(
1337 python: &dyn PythonSource,
1338 token: QueryToken<'_>,
1339 code_unit: &CodeUnit,
1340 raw: &str,
1341) -> Option<CodeUnit> {
1342 let trimmed = raw.trim();
1343 if trimmed.is_empty() {
1344 return None;
1345 }
1346
1347 let binder = python.import_binder_of(code_unit.source());
1348 if let Some((head, tail)) = trimmed.split_once('.') {
1349 if let Some(binding) = binder.bindings.get(head)
1350 && binding.kind == ImportKind::Namespace
1351 {
1352 let fq_name = format!("{}.{}", binding.module_specifier, tail);
1353 return python.definitions(&fq_name).next();
1354 }
1355 return python.definitions(trimmed).next();
1356 }
1357
1358 if let Some(binding) = binder.bindings.get(trimmed) {
1359 match binding.kind {
1360 ImportKind::Namespace => {
1361 return resolve_module_code_unit(
1362 python,
1363 Some(code_unit.source()),
1364 &binding.module_specifier,
1365 );
1366 }
1367 ImportKind::Named => {
1368 let imported_name = binding.imported_name.as_ref()?;
1369 let fqn = format!("{}.{}", binding.module_specifier, imported_name);
1370 return resolve_exported_fqn(python, Some(code_unit.source()), &fqn)
1371 .into_iter()
1372 .next()
1373 .or_else(|| python.definitions(&fqn).next());
1374 }
1375 _ => {}
1376 }
1377 }
1378
1379 if python
1380 .import_info_of(token, code_unit.source())
1381 .iter()
1382 .any(|import| import.is_wildcard)
1383 && let Some(imported) =
1384 resolve_import_bindings(python, token, code_unit.source()).get(trimmed)
1385 {
1386 return Some(imported.clone());
1387 }
1388
1389 if let Some(local) = resolve_preceding_local_class_base(python, token, code_unit, trimmed) {
1390 return Some(local);
1391 }
1392
1393 let local_fq_name = format!("{}.{}", code_unit.package_name(), trimmed);
1394 python
1395 .definitions(&local_fq_name)
1396 .next()
1397 .or_else(|| python.definitions(trimmed).next())
1398}
1399
1400fn class_definition_in_statement(node: Node<'_>) -> Option<Node<'_>> {
1401 if node.kind() == "class_definition" {
1402 return Some(node);
1403 }
1404 (node.kind() == "decorated_definition")
1405 .then(|| {
1406 let mut cursor = node.walk();
1407 node.named_children(&mut cursor)
1408 .find(|child| child.kind() == "class_definition")
1409 })
1410 .flatten()
1411}
1412
1413fn declaration_statement(node: Node<'_>) -> Option<Node<'_>> {
1414 match node.kind() {
1415 "decorated_definition" if class_definition_in_statement(node).is_some() => Some(node),
1416 "class_definition" => Some(
1417 node.parent()
1418 .filter(|parent| parent.kind() == "decorated_definition")
1419 .unwrap_or(node),
1420 ),
1421 _ => None,
1422 }
1423}
1424
1425fn scope_statement_binds_name(node: Node<'_>, source: &str, name: &str) -> bool {
1426 let mut stack = vec![node];
1427 while let Some(node) = stack.pop() {
1428 if node.kind() == "wildcard_import" {
1429 return true;
1430 }
1431 if matches!(node.kind(), "global_statement" | "nonlocal_statement") {
1432 let mut cursor = node.walk();
1433 if node.named_children(&mut cursor).any(|child| {
1434 child.kind() == "identifier"
1435 && child.utf8_text(source.as_bytes()).ok() == Some(name)
1436 }) {
1437 return true;
1438 }
1439 }
1440 let Some(bindings) = python_direct_scope_bindings_bounded(node, source, || true) else {
1441 return true;
1442 };
1443 if bindings
1444 .iter()
1445 .any(|binding| binding.declaration.utf8_text(source.as_bytes()).ok() == Some(name))
1446 {
1447 return true;
1448 }
1449 stack.extend(python_executable_scope_children(node));
1450 }
1451 false
1452}
1453
1454fn resolve_preceding_local_class_base(
1455 python: &dyn PythonSource,
1456 token: QueryToken<'_>,
1457 code_unit: &CodeUnit,
1458 raw: &str,
1459) -> Option<CodeUnit> {
1460 let parent = python.parent_of(code_unit)?;
1461 if !parent.is_function() && !parent.is_class() {
1462 return None;
1463 }
1464 let prepared = python.prepared_syntax(token, code_unit.source())?;
1465 let source = prepared.source();
1466 let current_node = prepared.declaration_node(code_unit)?;
1467 let current_statement = declaration_statement(current_node)?;
1468 let current_class = class_definition_in_statement(current_statement)?;
1469 let base_list = current_class.child_by_field_name("superclasses")?;
1470 let mut base_cursor = base_list.walk();
1471 if !base_list.named_children(&mut base_cursor).any(|base| {
1472 let origin = python_base_origin_node(base);
1473 origin.kind() == "identifier" && origin.utf8_text(source.as_bytes()).ok() == Some(raw)
1474 }) {
1475 return None;
1476 }
1477 let block = current_statement.parent()?;
1478 let mut block_cursor = block.walk();
1479 let siblings = block.named_children(&mut block_cursor).collect::<Vec<_>>();
1480 let current_index = siblings
1481 .iter()
1482 .position(|sibling| sibling.id() == current_statement.id())?;
1483
1484 let mut preceding: Option<(usize, CodeUnit, Node<'_>)> = None;
1485 for candidate in prepared.direct_children(&parent) {
1486 if !candidate.is_class()
1487 || candidate.source() != code_unit.source()
1488 || candidate.identifier() != raw
1489 {
1490 continue;
1491 }
1492 let Some(candidate_node) = prepared.declaration_node(candidate) else {
1493 continue;
1494 };
1495 let Some(candidate_statement) = declaration_statement(candidate_node) else {
1496 continue;
1497 };
1498 if candidate_statement.parent().map(|node| node.id()) != Some(block.id())
1499 || candidate_statement.start_byte() >= current_statement.start_byte()
1500 {
1501 continue;
1502 }
1503 if preceding
1504 .as_ref()
1505 .is_none_or(|(start, _, _)| candidate_statement.start_byte() > *start)
1506 {
1507 preceding = Some((
1508 candidate_statement.start_byte(),
1509 candidate.clone(),
1510 candidate_statement,
1511 ));
1512 }
1513 }
1514 let (_, candidate, candidate_statement) = preceding?;
1515 let candidate_index = siblings
1516 .iter()
1517 .position(|sibling| sibling.id() == candidate_statement.id())?;
1518 if candidate_index >= current_index
1519 || siblings[candidate_index + 1..current_index]
1520 .iter()
1521 .any(|statement| scope_statement_binds_name(*statement, source, raw))
1522 {
1523 return None;
1524 }
1525 Some(candidate)
1526}
1527
1528pub fn render_skeleton_recursive(
1529 index: &dyn CodeUnitIndex,
1530 code_unit: &CodeUnit,
1531 indent: &str,
1532 header_only: bool,
1533 out: &mut String,
1534) {
1535 enum Work {
1536 Declaration(CodeUnit, String),
1537 Elision(String),
1538 }
1539 let mut stack = vec![Work::Declaration(code_unit.clone(), indent.to_string())];
1540 while let Some(work) = stack.pop() {
1541 let Work::Declaration(unit, indent) = work else {
1542 let Work::Elision(indent) = work else {
1543 unreachable!()
1544 };
1545 out.push_str(&indent);
1546 out.push_str("[...]\n");
1547 continue;
1548 };
1549 if !unit.is_module()
1550 && !unit.is_file_scope()
1551 && let Some(signature) = index.signatures(&unit).first()
1552 {
1553 let signature = if unit.is_field() {
1556 signature.strip_suffix(" ...").unwrap_or(signature)
1557 } else {
1558 signature
1559 };
1560 for line in signature.lines() {
1561 out.push_str(&indent);
1562 out.push_str(line);
1563 out.push('\n');
1564 }
1565 }
1566 let children = index.direct_children(&unit);
1567 let child_indent = format!("{indent} ");
1568 if header_only && children.iter().any(|child| !child.is_field()) {
1569 stack.push(Work::Elision(child_indent.clone()));
1570 }
1571 for child in children
1572 .into_iter()
1573 .rev()
1574 .filter(|child| !header_only || child.is_field())
1575 {
1576 stack.push(Work::Declaration(child, child_indent.clone()));
1577 }
1578 }
1579}