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