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 visit: &mut dyn FnMut(&ProjectFile, Option<&dyn IndexedFileFacts>),
96 );
97}
98
99pub trait PythonUsageSource: PythonSource {
102 fn usage_index(&self) -> Arc<PythonUsageIndex>;
103}
104
105pub fn retain_modules_for_importer<T>(
124 project: &dyn Project,
125 importer: &ProjectFile,
126 candidates: &mut Vec<T>,
127 source_of: impl Fn(&T) -> &ProjectFile,
128) {
129 retain_nearest_owning_root(importer, candidates, |candidate| {
130 let source = source_of(candidate);
131 debug_assert_eq!(
132 importer.root(),
133 source.root(),
134 "one workspace's Python files share one project root"
135 );
136 python_import_root(project, source)
137 });
138}
139
140fn retain_nearest_owning_root<T>(
149 importer: &ProjectFile,
150 candidates: &mut Vec<T>,
151 root_of: impl Fn(&T) -> PathBuf,
152) {
153 if candidates.len() < 2 {
154 return;
155 }
156 let owning_depth = |candidate: &T| {
157 let root = root_of(candidate);
158 importer
159 .rel_path()
160 .starts_with(&root)
161 .then(|| root.components().count())
162 };
163 let depths = candidates.iter().map(owning_depth).collect::<Vec<_>>();
164 let Some(nearest) = depths.iter().copied().flatten().max() else {
165 return;
166 };
167 *candidates = std::mem::take(candidates)
168 .into_iter()
169 .zip(depths)
170 .filter(|(_, depth)| *depth == Some(nearest))
171 .map(|(candidate, _)| candidate)
172 .collect();
173}
174
175#[derive(Clone, Debug)]
177struct ProjectRootModuleName {
178 file: ProjectFile,
179 components: Vec<String>,
180}
181
182#[derive(Clone, Debug)]
191pub struct PythonSourceRootModule {
192 pub file: ProjectFile,
193 pub source_root: PathBuf,
197}
198
199#[derive(Default)]
221pub struct PythonModuleSpellings {
222 project: Option<Arc<dyn Project>>,
223 by_last_component: HashMap<String, Vec<ProjectRootModuleName>>,
224 package_dirs: HashSet<PathBuf>,
227 path_extensions: HashMap<PathBuf, HashSet<PathBuf>>,
239}
240
241impl std::fmt::Debug for PythonModuleSpellings {
242 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
243 f.debug_struct("PythonModuleSpellings")
244 .field("by_last_component", &self.by_last_component)
245 .field("package_dirs", &self.package_dirs)
246 .field("path_extensions", &self.path_extensions)
247 .finish_non_exhaustive()
248 }
249}
250
251impl PythonModuleSpellings {
252 pub(crate) fn project(&self) -> Option<&dyn Project> {
253 self.project.as_deref()
254 }
255
256 pub fn build(files: impl IntoIterator<Item = ProjectFile>, project: Arc<dyn Project>) -> Self {
262 let mut by_last_component: HashMap<String, Vec<ProjectRootModuleName>> = HashMap::default();
263 let mut package_dirs: HashSet<PathBuf> = HashSet::default();
264 let mut package_markers: Vec<ProjectFile> = Vec::new();
265 for file in files {
266 if file.rel_path().file_name().and_then(|name| name.to_str()) == Some("__init__.py")
267 && let Some(directory) = file.rel_path().parent()
268 {
269 package_dirs.insert(directory.to_path_buf());
270 package_markers.push(file.clone());
271 }
272 let components = python_project_module_components(&file);
273 let Some(last) = components.last() else {
274 continue;
275 };
276 by_last_component
277 .entry(last.clone())
278 .or_default()
279 .push(ProjectRootModuleName { file, components });
280 }
281 let mut path_extensions: HashMap<PathBuf, HashSet<PathBuf>> = HashMap::default();
282 for marker in package_markers {
283 let Some(directory) = marker.rel_path().parent() else {
284 continue;
285 };
286 if !python_package_holds_a_repeat_of_itself(directory, &package_dirs) {
287 continue;
288 }
289 let Ok(source) = project.read_source(&marker) else {
290 continue;
291 };
292 let extended = python_appended_path_dirs(&source, marker.rel_path());
293 if !extended.is_empty() {
294 path_extensions
295 .entry(directory.to_path_buf())
296 .or_default()
297 .extend(extended);
298 }
299 }
300 Self {
301 project: Some(project),
302 by_last_component,
303 package_dirs,
304 path_extensions,
305 }
306 }
307
308 pub fn suffix_candidates(&self, module_fq: &str) -> Vec<PythonSourceRootModule> {
320 if module_fq.starts_with('.') {
324 return Vec::new();
325 }
326 let components = parse_symbol_path(Language::Python, module_fq);
327 let Some(last) = components.last().filter(|_| components.len() > 1) else {
328 return Vec::new();
329 };
330 let mut modules = Vec::new();
331 for candidate in self
332 .by_last_component
333 .get(last.as_str())
334 .into_iter()
335 .flatten()
336 {
337 let Some(dropped) = candidate
340 .components
341 .len()
342 .checked_sub(components.len())
343 .filter(|dropped| *dropped > 0)
344 else {
345 continue;
346 };
347 if candidate.components[dropped..] != components[..] {
348 continue;
349 }
350 let source_root: PathBuf = candidate
353 .file
354 .rel_path()
355 .components()
356 .take(dropped)
357 .collect();
358 debug_assert!(
359 !source_root.as_os_str().is_empty(),
360 "dropping a component of {:?} leaves a directory",
361 candidate.file.rel_path()
362 );
363 modules.push(PythonSourceRootModule {
364 file: candidate.file.clone(),
365 source_root,
366 });
367 }
368 modules
369 }
370
371 pub fn establishes_source_root(&self, candidate: &PythonSourceRootModule) -> bool {
417 if let Some(declared) = python_declared_import_root(
418 self.project()
419 .expect("source-root candidates have a selected project"),
420 &candidate.file,
421 ) {
422 if declared == candidate.source_root {
423 return true;
424 }
425 let Some(extension) = python_repeated_root_extension_dir(
426 &candidate.file,
427 &candidate.source_root,
428 &declared,
429 ) else {
430 return false;
431 };
432 return self.package_chain_reaches_module(candidate)
433 && self
434 .path_extensions
435 .get(&candidate.source_root)
436 .is_some_and(|extended| extended.contains(&extension));
437 }
438 self.package_chain_reaches_module(candidate)
439 }
440
441 fn package_chain_reaches_module(&self, candidate: &PythonSourceRootModule) -> bool {
450 debug_assert!(
451 candidate
452 .file
453 .rel_path()
454 .starts_with(&candidate.source_root),
455 "a candidate's implied root is a prefix of its own path"
456 );
457 let mut directory = candidate.file.rel_path().parent();
458 while let Some(current) = directory {
459 if !self.package_dirs.contains(current) {
460 return false;
461 }
462 if current == candidate.source_root {
463 return true;
464 }
465 directory = current.parent();
466 }
467 false
468 }
469}
470
471fn python_package_holds_a_repeat_of_itself(
479 directory: &Path,
480 package_dirs: &HashSet<PathBuf>,
481) -> bool {
482 directory
483 .file_name()
484 .is_some_and(|name| package_dirs.contains(&directory.join(name)))
485}
486
487fn python_appended_path_dirs(source: &str, marker_rel_path: &Path) -> Vec<PathBuf> {
502 let Some(tree) = parse_python_tree(source) else {
503 return Vec::new();
504 };
505 let mut bound: HashMap<String, PathBuf> = HashMap::default();
506 bound.insert("__file__".to_owned(), marker_rel_path.to_path_buf());
507 let mut appended: Vec<PathBuf> = Vec::new();
508 walk_named_tree_preorder(tree.root_node(), true, |node| {
509 match node.kind() {
510 "function_definition" | "class_definition" | "lambda" => {
514 return WalkControl::SkipChildren;
515 }
516 "assignment" => {
517 if python_names_path(node.child_by_field_name("left"), source) {
522 appended.clear();
523 }
524 python_bind_assignment(node, source, &mut bound);
525 }
526 "call" => {
527 let directories = python_path_appended_dirs_in_call(node, source, &bound);
528 if directories.is_empty()
529 || !python_extension_statement_runs(node, source, &bound, &directories)
530 {
531 return WalkControl::Continue;
532 }
533 for directory in directories {
534 if directory.is_relative()
535 && !directory.as_os_str().is_empty()
536 && !appended.contains(&directory)
537 {
538 appended.push(directory);
539 }
540 }
541 }
542 _ => {}
543 }
544 WalkControl::Continue
545 });
546 appended
547}
548
549fn python_names_path(node: Option<Node<'_>>, source: &str) -> bool {
551 node.is_some_and(|node| {
552 node.kind() == "identifier" && node.utf8_text(source.as_bytes()).ok() == Some("__path__")
553 })
554}
555
556fn python_extension_statement_runs(
571 call: Node<'_>,
572 source: &str,
573 bound: &HashMap<String, PathBuf>,
574 directories: &[PathBuf],
575) -> bool {
576 let Some(statement) = call.parent() else {
577 return false;
578 };
579 if statement.kind() != "expression_statement" {
580 return false;
581 }
582 let Some(container) = statement.parent() else {
583 return false;
584 };
585 if container.kind() == "module" {
586 return true;
587 }
588 let guarded = if container.kind() == "block" {
593 container.parent().filter(|parent| {
594 parent.kind() == "if_statement"
595 && parent.child_by_field_name("consequence") == Some(container)
596 })
597 } else if container.kind() == "if_statement"
598 && container.child_by_field_name("consequence") == Some(statement)
599 {
600 Some(container)
601 } else {
602 None
603 };
604 let Some(guarded) = guarded else {
605 return false;
606 };
607 if guarded.parent().map(|parent| parent.kind()) != Some("module") {
612 return false;
613 }
614 let Some(condition) = guarded.child_by_field_name("condition") else {
615 return false;
616 };
617 python_condition_justifies_path_extension(condition, source, bound, directories)
618}
619
620fn python_condition_justifies_path_extension(
625 condition: Node<'_>,
626 source: &str,
627 bound: &HashMap<String, PathBuf>,
628 directories: &[PathBuf],
629) -> bool {
630 if condition.kind() != "comparison_operator" {
631 return false;
632 }
633 let mut cursor = condition.walk();
634 let operators = condition
635 .children_by_field_name("operators", &mut cursor)
636 .map(|operator| operator.kind())
637 .collect::<Vec<_>>();
638 if operators.as_slice() != ["not in"] {
639 return false;
640 }
641 let mut cursor = condition.walk();
644 let operands: Vec<Node<'_>> = condition.named_children(&mut cursor).collect();
645 let [left, right] = operands.as_slice() else {
646 return false;
647 };
648 if !python_names_path(Some(*right), source) {
649 return false;
650 }
651 python_path_expression_dir(*left, source, bound)
652 .is_some_and(|directory| directories.contains(&directory))
653}
654
655fn python_bind_assignment(
660 assignment: Node<'_>,
661 source: &str,
662 bound: &mut HashMap<String, PathBuf>,
663) {
664 let mut targets: Vec<Node<'_>> = Vec::new();
665 let mut current = assignment;
666 loop {
667 let (Some(left), Some(right)) = (
668 current.child_by_field_name("left"),
669 current.child_by_field_name("right"),
670 ) else {
671 return;
672 };
673 targets.push(left);
674 if right.kind() != "assignment" {
675 let resolved = python_path_expression_dir(right, source, bound);
676 for target in targets {
677 if target.kind() != "identifier" {
678 continue;
679 }
680 let Ok(name) = target.utf8_text(source.as_bytes()) else {
681 continue;
682 };
683 match &resolved {
684 Some(directory) => {
685 bound.insert(name.to_owned(), directory.clone());
686 }
687 None => {
688 bound.remove(name);
689 }
690 }
691 }
692 return;
693 }
694 current = right;
695 }
696}
697
698fn python_path_appended_dirs_in_call(
705 call: Node<'_>,
706 source: &str,
707 bound: &HashMap<String, PathBuf>,
708) -> Vec<PathBuf> {
709 let Some(function) = call.child_by_field_name("function") else {
710 return Vec::new();
711 };
712 let Some(path) = python_static_attribute_path(function) else {
713 return Vec::new();
714 };
715 let [object, method] = path.as_slice() else {
716 return Vec::new();
717 };
718 if object.utf8_text(source.as_bytes()).ok() != Some("__path__") {
719 return Vec::new();
720 }
721 let Ok(method) = method.utf8_text(source.as_bytes()) else {
722 return Vec::new();
723 };
724 let Some(arguments) = call.child_by_field_name("arguments") else {
725 return Vec::new();
726 };
727 let mut cursor = arguments.walk();
728 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
729 let named: Vec<Node<'_>> = match method {
730 "append" => arguments.first().copied().into_iter().collect(),
731 "insert" => arguments.get(1).copied().into_iter().collect(),
732 "extend" => match arguments.first() {
733 Some(list) if list.kind() == "list" => {
734 let mut cursor = list.walk();
735 list.named_children(&mut cursor).collect()
736 }
737 Some(iterator) => vec![*iterator],
738 None => Vec::new(),
739 },
740 _ => Vec::new(),
741 };
742 named
743 .into_iter()
744 .filter_map(|node| python_path_expression_dir(node, source, bound))
745 .collect()
746}
747
748fn python_path_expression_dir(
752 node: Node<'_>,
753 source: &str,
754 bound: &HashMap<String, PathBuf>,
755) -> Option<PathBuf> {
756 if let Some(literal) = python_plain_string_literal(node, source) {
757 let directory = PathBuf::from(literal);
758 return (!directory.as_os_str().is_empty()).then_some(directory);
759 }
760 match node.kind() {
761 "identifier" => bound.get(node.utf8_text(source.as_bytes()).ok()?).cloned(),
762 "call" => python_path_call_dir(node, source, bound),
763 _ => None,
764 }
765}
766
767fn python_path_call_dir(
770 call: Node<'_>,
771 source: &str,
772 bound: &HashMap<String, PathBuf>,
773) -> Option<PathBuf> {
774 let path = python_static_attribute_path(call.child_by_field_name("function")?)?;
775 let names: Vec<&str> = path
776 .iter()
777 .map(|node| node.utf8_text(source.as_bytes()).ok())
778 .collect::<Option<Vec<_>>>()?;
779 let arguments = call.child_by_field_name("arguments")?;
780 let mut cursor = arguments.walk();
781 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
782 match names.as_slice() {
783 ["os", "path", "dirname"] => {
784 let directory = python_path_expression_dir(*arguments.first()?, source, bound)?;
785 Some(directory.parent()?.to_path_buf())
786 }
787 ["os", "path", "join"] => {
788 let mut parts = arguments.iter();
789 let mut joined = python_path_expression_dir(*parts.next()?, source, bound)?;
790 for part in parts {
791 joined.push(python_path_expression_dir(*part, source, bound)?);
792 }
793 Some(joined)
794 }
795 _ => None,
796 }
797}
798
799pub fn resolve_module_code_unit(
807 python: &dyn PythonSource,
808 importer: Option<&ProjectFile>,
809 module_fq: &str,
810) -> Option<CodeUnit> {
811 let indexed = match python.path_module_fqn(module_fq) {
812 Some(units) => indexed_module_for_importer(python.project(), importer, units),
815 None => indexed_module_for_importer(
816 python.project(),
817 importer,
818 python.definition_fqn(module_fq),
819 ),
820 };
821 match indexed {
822 IndexedModuleLookup::Unique(module) => Some(module),
823 IndexedModuleLookup::Ambiguous => None,
827 IndexedModuleLookup::Missing => {
828 let importer = importer?;
829 module_under_source_root(&python.module_spellings(), importer, module_fq)
830 }
831 }
832}
833
834enum IndexedModuleLookup {
843 Unique(CodeUnit),
844 Ambiguous,
845 Missing,
846}
847
848fn indexed_module_for_importer(
849 project: &dyn Project,
850 importer: Option<&ProjectFile>,
851 units: Vec<CodeUnit>,
852) -> IndexedModuleLookup {
853 let mut modules = units
854 .into_iter()
855 .filter(CodeUnit::is_module)
856 .collect::<Vec<_>>();
857 if let Some(importer) = importer {
858 retain_modules_for_importer(project, importer, &mut modules, CodeUnit::source);
859 }
860 match modules.len() {
861 0 => IndexedModuleLookup::Missing,
862 1 => IndexedModuleLookup::Unique(modules.remove(0)),
863 _ => IndexedModuleLookup::Ambiguous,
864 }
865}
866
867pub fn python_module_file_under_source_root(
896 spellings: &PythonModuleSpellings,
897 importer: &ProjectFile,
898 module_fq: &str,
899) -> Option<ProjectFile> {
900 let mut candidates = spellings.suffix_candidates(module_fq);
901 retain_nearest_owning_root(importer, &mut candidates, |candidate| {
902 candidate.source_root.clone()
903 });
904 candidates.retain(|candidate| spellings.establishes_source_root(candidate));
905 let [candidate] = candidates.as_slice() else {
906 return None;
907 };
908 debug_assert_eq!(
909 importer.root(),
910 candidate.file.root(),
911 "one workspace's Python files share one project root"
912 );
913 Some(candidate.file.clone())
914}
915
916fn module_under_source_root(
921 spellings: &PythonModuleSpellings,
922 importer: &ProjectFile,
923 module_fq: &str,
924) -> Option<CodeUnit> {
925 let file = python_module_file_under_source_root(spellings, importer, module_fq)?;
926 module_code_unit(spellings.project()?, &file)
927}
928
929pub fn resolve_module_code_units_batch(
935 python: &dyn PythonSource,
936 importer: Option<&ProjectFile>,
937 module_fqs: &[String],
938) -> Vec<Option<CodeUnit>> {
939 let path_results = python.path_module_fqns_batch(module_fqs);
940 let mut lookups: Vec<IndexedModuleLookup> = Vec::with_capacity(module_fqs.len());
941 let mut needs_definition_fallback = Vec::new();
942 for (i, units) in path_results.into_iter().enumerate() {
943 match units {
944 Some(units) => lookups.push(indexed_module_for_importer(
945 python.project(),
946 importer,
947 units,
948 )),
949 None => {
950 needs_definition_fallback.push(i);
951 lookups.push(IndexedModuleLookup::Missing);
953 }
954 }
955 }
956 for i in needs_definition_fallback {
957 lookups[i] = indexed_module_for_importer(
958 python.project(),
959 importer,
960 python.definition_fqn(&module_fqs[i]),
961 );
962 }
963 let mut results: Vec<Option<CodeUnit>> = vec![None; module_fqs.len()];
964 let mut needs_source_root = Vec::new();
965 for (i, lookup) in lookups.into_iter().enumerate() {
966 match lookup {
967 IndexedModuleLookup::Unique(module) => results[i] = Some(module),
968 IndexedModuleLookup::Ambiguous => {}
969 IndexedModuleLookup::Missing => needs_source_root.push(i),
970 }
971 }
972 if let Some(importer) = importer
973 && !needs_source_root.is_empty()
974 {
975 let spellings = python.module_spellings();
978 for i in needs_source_root {
979 results[i] = module_under_source_root(&spellings, importer, &module_fqs[i]);
980 }
981 }
982 results
983}
984
985pub fn compute_export_index_of(
986 python: &dyn PythonSource,
987 token: QueryToken<'_>,
988 file: &ProjectFile,
989) -> ExportIndex {
990 let mut index = ExportIndex::empty();
991 let mut events = Vec::new();
992 let declarations = python.top_level_declarations(file);
993 collect_local_export_events(
994 declarations.iter(),
995 |code_unit| {
996 python
997 .ranges(code_unit)
998 .iter()
999 .map(|range| range.start_byte)
1000 .max()
1001 .unwrap_or(usize::MAX)
1002 },
1003 &mut events,
1004 );
1005
1006 let imports = python.import_info_of(token, file);
1007 collect_reexport_events_from_imports(python, file, &imports, &mut events, &mut index);
1008
1009 finish_export_index(events, index)
1010}
1011
1012pub fn export_index_from_file_facts(
1013 python: &dyn PythonSource,
1014 file: &ProjectFile,
1015 facts: &dyn IndexedFileFacts,
1016 module_name: &str,
1017) -> ExportIndex {
1018 let mut index = ExportIndex::empty();
1019 let mut events = Vec::new();
1020 collect_local_export_events(
1021 facts.top_level_declarations().iter(),
1022 |code_unit| {
1023 facts
1024 .declaration_ranges(code_unit)
1025 .into_iter()
1026 .flatten()
1027 .map(|range| range.start_byte)
1028 .max()
1029 .unwrap_or(usize::MAX)
1030 },
1031 &mut events,
1032 );
1033
1034 if !facts
1035 .top_level_declarations()
1036 .iter()
1037 .any(CodeUnit::is_module)
1038 && let Some(identifier) = module_name.rsplit('.').next()
1039 && !identifier.is_empty()
1040 && !identifier.starts_with('_')
1041 {
1042 events.push((
1043 0,
1044 identifier.to_string(),
1045 ExportEntry::Local {
1046 local_name: identifier.to_string(),
1047 },
1048 ));
1049 }
1050
1051 collect_reexport_events_from_imports(python, file, facts.imports(), &mut events, &mut index);
1052
1053 finish_export_index(events, index)
1054}
1055
1056fn collect_local_export_events<'a>(
1057 declarations: impl IntoIterator<Item = &'a CodeUnit>,
1058 mut start_byte: impl FnMut(&CodeUnit) -> usize,
1059 events: &mut Vec<(usize, String, ExportEntry)>,
1060) -> HashSet<String> {
1061 let mut local_names = HashSet::default();
1062 for code_unit in declarations {
1063 let identifier = code_unit.identifier().trim();
1064 if identifier.is_empty() {
1065 continue;
1066 }
1067 local_names.insert(identifier.to_string());
1068 events.push((
1069 start_byte(code_unit),
1070 identifier.to_string(),
1071 ExportEntry::Local {
1072 local_name: identifier.to_string(),
1073 },
1074 ));
1075 }
1076 local_names
1077}
1078
1079fn finish_export_index(
1080 mut events: Vec<(usize, String, ExportEntry)>,
1081 mut index: ExportIndex,
1082) -> ExportIndex {
1083 events.sort_by_key(|(start_byte, _, _)| *start_byte);
1084 for (_, exported_name, entry) in events {
1085 index.exports_by_name.insert(exported_name, entry);
1086 }
1087 index
1088}
1089
1090fn collect_reexport_events_from_imports(
1091 python: &dyn PythonSource,
1092 file: &ProjectFile,
1093 imports: &[ImportInfo],
1094 events: &mut Vec<(usize, String, ExportEntry)>,
1095 index: &mut ExportIndex,
1096) {
1097 for import in imports {
1098 if import
1099 .path
1100 .as_ref()
1101 .is_none_or(|path| !path.lexical_scopes.is_empty())
1102 {
1103 continue;
1104 }
1105 record_single_reexport_event(python, file, import, events, index);
1106 }
1107}
1108
1109fn record_single_reexport_event(
1110 python: &dyn PythonSource,
1111 file: &ProjectFile,
1112 import: &ImportInfo,
1113 events: &mut Vec<(usize, String, ExportEntry)>,
1114 index: &mut ExportIndex,
1115) {
1116 let Some(PythonImportDetails::FromImport {
1117 module,
1118 name,
1119 alias,
1120 wildcard,
1121 }) = python_import_details(import)
1122 else {
1123 return;
1124 };
1125 let start_byte = import
1126 .path
1127 .as_ref()
1128 .map(|path| path.declaration_start_byte)
1129 .unwrap_or(usize::MAX);
1130 let resolved_module = if module.starts_with('.') {
1131 resolve_python_relative_module(python.project(), file, &module)
1132 } else {
1133 Some(module.clone())
1134 };
1135 let Some(resolved_module) = resolved_module else {
1136 return;
1137 };
1138
1139 if wildcard {
1140 index.reexport_stars.push(ReexportStar {
1141 module_specifier: resolved_module,
1142 });
1143 return;
1144 }
1145 let exported_name = alias.unwrap_or(name.clone());
1146 let module_candidate = format!("{resolved_module}.{name}");
1152 if resolve_module_code_unit(python, Some(file), &module_candidate).is_some() {
1153 events.push((
1154 start_byte,
1155 exported_name,
1156 ExportEntry::ReexportedModule {
1157 module_specifier: module_candidate,
1158 },
1159 ));
1160 return;
1161 }
1162 events.push((
1163 start_byte,
1164 exported_name,
1165 ExportEntry::ReexportedNamed {
1166 module_specifier: resolved_module,
1167 imported_name: name,
1168 },
1169 ));
1170}
1171
1172pub fn import_binder_from_imports(
1173 python: &dyn PythonSource,
1174 file: &ProjectFile,
1175 imports: &[ImportInfo],
1176) -> ImportBinder {
1177 let mut binder = ImportBinder::empty();
1178
1179 for (local_name, binding) in import_bindings_from_imports(python, file, imports) {
1180 binder.bindings.insert(local_name, binding);
1181 }
1182
1183 binder
1184}
1185
1186pub fn import_bindings_from_imports(
1191 python: &dyn PythonSource,
1192 file: &ProjectFile,
1193 imports: &[ImportInfo],
1194) -> Vec<(String, ImportBinding)> {
1195 let mut bindings = Vec::new();
1196
1197 for import in imports {
1198 let Some(details) = python_import_details(import) else {
1199 continue;
1200 };
1201 match details {
1202 PythonImportDetails::Import { module, alias } => {
1203 let local_name = python_namespace_binding_name(import, alias.as_deref(), &module);
1204 let module_specifier =
1205 python_namespace_binding_module(import, alias.as_deref(), &module);
1206 bindings.push((
1207 local_name,
1208 ImportBinding {
1209 module_specifier,
1210 namespace_imported_module: Some(module),
1211 kind: ImportKind::Namespace,
1212 imported_name: None,
1213 },
1214 ));
1215 }
1216 PythonImportDetails::FromImport {
1217 module,
1218 name,
1219 wildcard,
1220 ..
1221 } => {
1222 let resolved_module = if module.starts_with('.') {
1223 resolve_python_relative_module(python.project(), file, &module)
1224 } else {
1225 Some(module.clone())
1226 };
1227 let Some(resolved_module) = resolved_module else {
1228 continue;
1229 };
1230 if wildcard {
1231 bindings.extend(
1236 public_declarations_in_module(python, file, &resolved_module)
1237 .into_iter()
1238 .map(|declaration| {
1239 let name = declaration.identifier().to_string();
1240 (
1241 name.clone(),
1242 ImportBinding {
1243 module_specifier: resolved_module.clone(),
1244 namespace_imported_module: None,
1245 kind: ImportKind::Named,
1246 imported_name: Some(name),
1247 },
1248 )
1249 }),
1250 );
1251 continue;
1252 }
1253 let local_name = import
1258 .local_name()
1259 .map(str::to_string)
1260 .unwrap_or_else(|| name.clone());
1261 let module_candidate = format!("{resolved_module}.{name}");
1262 let exported = resolve_exported_fqn(python, Some(file), &module_candidate);
1263 let exported_module =
1264 if !exported.is_empty() && exported.iter().all(CodeUnit::is_module) {
1265 let mut paths = exported.iter().map(CodeUnit::fq_segment_texts);
1266 let path = paths
1267 .next()
1268 .expect("non-empty exported modules have a first path");
1269 paths
1270 .all(|candidate| candidate == path)
1271 .then(|| path.join("."))
1272 } else {
1273 None
1274 };
1275 let implicit_submodule = if exported.is_empty()
1276 && resolve_module_code_unit(python, Some(file), &module_candidate).is_some()
1277 {
1278 Some(module_candidate)
1279 } else {
1280 None
1281 };
1282 if let Some(module_specifier) = exported_module.or(implicit_submodule) {
1283 bindings.push((
1284 local_name,
1285 ImportBinding {
1286 module_specifier,
1287 namespace_imported_module: None,
1288 kind: ImportKind::Namespace,
1289 imported_name: None,
1290 },
1291 ));
1292 continue;
1293 }
1294 bindings.push((
1295 local_name,
1296 ImportBinding {
1297 module_specifier: resolved_module,
1298 namespace_imported_module: None,
1299 kind: ImportKind::Named,
1300 imported_name: Some(name),
1301 },
1302 ));
1303 }
1304 }
1305 }
1306
1307 bindings
1308}
1309
1310pub fn public_declarations_in_module(
1311 python: &dyn PythonSource,
1312 importer: &ProjectFile,
1313 module_fq: &str,
1314) -> Vec<CodeUnit> {
1315 let Some(module_code_unit) = resolve_module_code_unit(python, Some(importer), module_fq) else {
1316 return Vec::new();
1317 };
1318 python
1319 .direct_children(&module_code_unit)
1320 .into_iter()
1321 .filter(|code_unit| !code_unit.identifier().starts_with('_'))
1322 .collect()
1323}
1324
1325pub fn resolve_base_class(
1326 python: &dyn PythonSource,
1327 token: QueryToken<'_>,
1328 code_unit: &CodeUnit,
1329 raw: &str,
1330) -> Option<CodeUnit> {
1331 let trimmed = raw.trim();
1332 if trimmed.is_empty() {
1333 return None;
1334 }
1335
1336 let binder = python.import_binder_of(code_unit.source());
1337 if let Some((head, tail)) = trimmed.split_once('.') {
1338 if let Some(binding) = binder.bindings.get(head)
1339 && binding.kind == ImportKind::Namespace
1340 {
1341 let fq_name = format!("{}.{}", binding.module_specifier, tail);
1342 return python.definitions(&fq_name).next();
1343 }
1344 return python.definitions(trimmed).next();
1345 }
1346
1347 if let Some(binding) = binder.bindings.get(trimmed) {
1348 match binding.kind {
1349 ImportKind::Namespace => {
1350 return resolve_module_code_unit(
1351 python,
1352 Some(code_unit.source()),
1353 &binding.module_specifier,
1354 );
1355 }
1356 ImportKind::Named => {
1357 let imported_name = binding.imported_name.as_ref()?;
1358 let fqn = format!("{}.{}", binding.module_specifier, imported_name);
1359 return resolve_exported_fqn(python, Some(code_unit.source()), &fqn)
1360 .into_iter()
1361 .next()
1362 .or_else(|| python.definitions(&fqn).next());
1363 }
1364 _ => {}
1365 }
1366 }
1367
1368 if python
1369 .import_info_of(token, code_unit.source())
1370 .iter()
1371 .any(|import| import.is_wildcard)
1372 && let Some(imported) =
1373 resolve_import_bindings(python, token, code_unit.source()).get(trimmed)
1374 {
1375 return Some(imported.clone());
1376 }
1377
1378 if let Some(local) = resolve_preceding_local_class_base(python, token, code_unit, trimmed) {
1379 return Some(local);
1380 }
1381
1382 let local_fq_name = format!("{}.{}", code_unit.package_name(), trimmed);
1383 python
1384 .definitions(&local_fq_name)
1385 .next()
1386 .or_else(|| python.definitions(trimmed).next())
1387}
1388
1389fn class_definition_in_statement(node: Node<'_>) -> Option<Node<'_>> {
1390 if node.kind() == "class_definition" {
1391 return Some(node);
1392 }
1393 (node.kind() == "decorated_definition")
1394 .then(|| {
1395 let mut cursor = node.walk();
1396 node.named_children(&mut cursor)
1397 .find(|child| child.kind() == "class_definition")
1398 })
1399 .flatten()
1400}
1401
1402fn declaration_statement(node: Node<'_>) -> Option<Node<'_>> {
1403 match node.kind() {
1404 "decorated_definition" if class_definition_in_statement(node).is_some() => Some(node),
1405 "class_definition" => Some(
1406 node.parent()
1407 .filter(|parent| parent.kind() == "decorated_definition")
1408 .unwrap_or(node),
1409 ),
1410 _ => None,
1411 }
1412}
1413
1414fn scope_statement_binds_name(node: Node<'_>, source: &str, name: &str) -> bool {
1415 let mut stack = vec![node];
1416 while let Some(node) = stack.pop() {
1417 if node.kind() == "wildcard_import" {
1418 return true;
1419 }
1420 if matches!(node.kind(), "global_statement" | "nonlocal_statement") {
1421 let mut cursor = node.walk();
1422 if node.named_children(&mut cursor).any(|child| {
1423 child.kind() == "identifier"
1424 && child.utf8_text(source.as_bytes()).ok() == Some(name)
1425 }) {
1426 return true;
1427 }
1428 }
1429 let Some(bindings) = python_direct_scope_bindings_bounded(node, source, || true) else {
1430 return true;
1431 };
1432 if bindings
1433 .iter()
1434 .any(|binding| binding.declaration.utf8_text(source.as_bytes()).ok() == Some(name))
1435 {
1436 return true;
1437 }
1438 stack.extend(python_executable_scope_children(node));
1439 }
1440 false
1441}
1442
1443fn resolve_preceding_local_class_base(
1444 python: &dyn PythonSource,
1445 token: QueryToken<'_>,
1446 code_unit: &CodeUnit,
1447 raw: &str,
1448) -> Option<CodeUnit> {
1449 let parent = python.parent_of(code_unit)?;
1450 if !parent.is_function() && !parent.is_class() {
1451 return None;
1452 }
1453 let prepared = python.prepared_syntax(token, code_unit.source())?;
1454 let source = prepared.source();
1455 let current_node = prepared.declaration_node(code_unit)?;
1456 let current_statement = declaration_statement(current_node)?;
1457 let current_class = class_definition_in_statement(current_statement)?;
1458 let base_list = current_class.child_by_field_name("superclasses")?;
1459 let mut base_cursor = base_list.walk();
1460 if !base_list.named_children(&mut base_cursor).any(|base| {
1461 let origin = python_base_origin_node(base);
1462 origin.kind() == "identifier" && origin.utf8_text(source.as_bytes()).ok() == Some(raw)
1463 }) {
1464 return None;
1465 }
1466 let block = current_statement.parent()?;
1467 let mut block_cursor = block.walk();
1468 let siblings = block.named_children(&mut block_cursor).collect::<Vec<_>>();
1469 let current_index = siblings
1470 .iter()
1471 .position(|sibling| sibling.id() == current_statement.id())?;
1472
1473 let mut preceding: Option<(usize, CodeUnit, Node<'_>)> = None;
1474 for candidate in prepared.direct_children(&parent) {
1475 if !candidate.is_class()
1476 || candidate.source() != code_unit.source()
1477 || candidate.identifier() != raw
1478 {
1479 continue;
1480 }
1481 let Some(candidate_node) = prepared.declaration_node(candidate) else {
1482 continue;
1483 };
1484 let Some(candidate_statement) = declaration_statement(candidate_node) else {
1485 continue;
1486 };
1487 if candidate_statement.parent().map(|node| node.id()) != Some(block.id())
1488 || candidate_statement.start_byte() >= current_statement.start_byte()
1489 {
1490 continue;
1491 }
1492 if preceding
1493 .as_ref()
1494 .is_none_or(|(start, _, _)| candidate_statement.start_byte() > *start)
1495 {
1496 preceding = Some((
1497 candidate_statement.start_byte(),
1498 candidate.clone(),
1499 candidate_statement,
1500 ));
1501 }
1502 }
1503 let (_, candidate, candidate_statement) = preceding?;
1504 let candidate_index = siblings
1505 .iter()
1506 .position(|sibling| sibling.id() == candidate_statement.id())?;
1507 if candidate_index >= current_index
1508 || siblings[candidate_index + 1..current_index]
1509 .iter()
1510 .any(|statement| scope_statement_binds_name(*statement, source, raw))
1511 {
1512 return None;
1513 }
1514 Some(candidate)
1515}
1516
1517pub fn render_skeleton_recursive(
1518 index: &dyn CodeUnitIndex,
1519 code_unit: &CodeUnit,
1520 indent: &str,
1521 header_only: bool,
1522 out: &mut String,
1523) {
1524 enum Work {
1525 Declaration(CodeUnit, String),
1526 Elision(String),
1527 }
1528 let mut stack = vec![Work::Declaration(code_unit.clone(), indent.to_string())];
1529 while let Some(work) = stack.pop() {
1530 let Work::Declaration(unit, indent) = work else {
1531 let Work::Elision(indent) = work else {
1532 unreachable!()
1533 };
1534 out.push_str(&indent);
1535 out.push_str("[...]\n");
1536 continue;
1537 };
1538 if !unit.is_module()
1539 && !unit.is_file_scope()
1540 && let Some(signature) = index.signatures(&unit).first()
1541 {
1542 let signature = if unit.is_field() {
1545 signature.strip_suffix(" ...").unwrap_or(signature)
1546 } else {
1547 signature
1548 };
1549 for line in signature.lines() {
1550 out.push_str(&indent);
1551 out.push_str(line);
1552 out.push('\n');
1553 }
1554 }
1555 let children = index.direct_children(&unit);
1556 let child_indent = format!("{indent} ");
1557 if header_only && children.iter().any(|child| !child.is_field()) {
1558 stack.push(Work::Elision(child_indent.clone()));
1559 }
1560 for child in children
1561 .into_iter()
1562 .rev()
1563 .filter(|child| !header_only || child.is_field())
1564 {
1565 stack.push(Work::Declaration(child, child_indent.clone()));
1566 }
1567 }
1568}