1use brokk_bifrost_core::analyzer::capabilities::ImportAnalysisProvider;
13use brokk_bifrost_core::analyzer::model::{CodeUnitType, ImportInfo};
14use brokk_bifrost_core::analyzer::prepared_syntax::{IndexedFileFacts, PreparedSyntaxTree};
15use brokk_bifrost_core::analyzer::query_token::QueryToken;
16use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
17use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
18use brokk_bifrost_core::analyzer::usages::model::{
19 ExportEntry, ExportIndex, ImportBinder, ImportBinding, ImportKind, ReexportStar,
20};
21use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, Language, ProjectFile};
22use brokk_bifrost_core::hash::{HashMap, HashSet};
23use std::collections::BTreeSet;
24use std::path::{Path, PathBuf};
25use std::sync::Arc;
26use tree_sitter::Node;
27
28use crate::declarations::{
29 collect_python_identifiers, module_code_unit, parse_python_tree, python_declared_import_root,
30 python_import_root, python_module_name, python_project_module_components,
31 python_repeated_root_extension_dir,
32};
33use crate::imports::{
34 PythonImportDetails, python_import_details, python_import_infos_from_node,
35 python_namespace_binding_module, python_namespace_binding_name, resolve_exported_fqn,
36 resolve_import_bindings, 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: CodeUnitIndex + ImportAnalysisProvider {
52 fn path_module_fqn(&self, module_fq: &str) -> Option<Vec<CodeUnit>>;
55
56 fn path_module_fqns_batch(&self, module_fqs: &[String]) -> Vec<Option<Vec<CodeUnit>>>;
59
60 fn definition_fqn(&self, fqn: &str) -> Vec<CodeUnit>;
61
62 fn module_spellings(&self) -> Arc<PythonModuleSpellings>;
69
70 fn import_binder_of(&self, file: &ProjectFile) -> Arc<ImportBinder>;
75
76 fn export_index_of(&self, file: &ProjectFile) -> Arc<ExportIndex>;
77
78 fn prepared_syntax(
90 &self,
91 token: QueryToken<'_>,
92 file: &ProjectFile,
93 ) -> Option<Arc<PreparedSyntaxTree>>;
94
95 fn visit_file_facts(
98 &self,
99 files: &[ProjectFile],
100 visit: &mut dyn FnMut(&ProjectFile, Option<&dyn IndexedFileFacts>),
101 );
102}
103
104pub trait PythonUsageSource: PythonSource {
107 fn usage_index(&self) -> Arc<PythonUsageIndex>;
108}
109
110pub fn extract_type_identifiers(source: &str) -> BTreeSet<String> {
111 let Some(tree) = parse_python_tree(source) else {
112 return BTreeSet::new();
113 };
114 let mut identifiers = HashSet::default();
115 collect_python_identifiers(tree.root_node(), source, &mut identifiers);
116 identifiers.into_iter().collect()
117}
118
119pub fn retain_modules_for_importer<T>(
138 importer: &ProjectFile,
139 candidates: &mut Vec<T>,
140 source_of: impl Fn(&T) -> &ProjectFile,
141) {
142 retain_nearest_owning_root(importer, candidates, |candidate| {
143 let source = source_of(candidate);
144 debug_assert_eq!(
145 importer.root(),
146 source.root(),
147 "one workspace's Python files share one project root"
148 );
149 python_import_root(source)
150 });
151}
152
153fn retain_nearest_owning_root<T>(
162 importer: &ProjectFile,
163 candidates: &mut Vec<T>,
164 root_of: impl Fn(&T) -> PathBuf,
165) {
166 if candidates.len() < 2 {
167 return;
168 }
169 let owning_depth = |candidate: &T| {
170 let root = root_of(candidate);
171 importer
172 .rel_path()
173 .starts_with(&root)
174 .then(|| root.components().count())
175 };
176 let depths = candidates.iter().map(owning_depth).collect::<Vec<_>>();
177 let Some(nearest) = depths.iter().copied().flatten().max() else {
178 return;
179 };
180 *candidates = std::mem::take(candidates)
181 .into_iter()
182 .zip(depths)
183 .filter(|(_, depth)| *depth == Some(nearest))
184 .map(|(candidate, _)| candidate)
185 .collect();
186}
187
188#[derive(Clone, Debug)]
190struct ProjectRootModuleName {
191 file: ProjectFile,
192 components: Vec<String>,
193}
194
195#[derive(Clone, Debug)]
204pub struct PythonSourceRootModule {
205 pub file: ProjectFile,
206 pub source_root: PathBuf,
210}
211
212#[derive(Debug, Default)]
234pub struct PythonModuleSpellings {
235 by_last_component: HashMap<String, Vec<ProjectRootModuleName>>,
236 package_dirs: HashSet<PathBuf>,
239 path_extensions: HashMap<PathBuf, HashSet<PathBuf>>,
251}
252
253impl PythonModuleSpellings {
254 pub fn build(
261 files: impl IntoIterator<Item = ProjectFile>,
262 mut read_source: impl FnMut(&ProjectFile) -> Option<String>,
263 ) -> Self {
264 let mut by_last_component: HashMap<String, Vec<ProjectRootModuleName>> = HashMap::default();
265 let mut package_dirs: HashSet<PathBuf> = HashSet::default();
266 let mut package_markers: Vec<ProjectFile> = Vec::new();
267 for file in files {
268 if file.rel_path().file_name().and_then(|name| name.to_str()) == Some("__init__.py")
269 && let Some(directory) = file.rel_path().parent()
270 {
271 package_dirs.insert(directory.to_path_buf());
272 package_markers.push(file.clone());
273 }
274 let components = python_project_module_components(&file);
275 let Some(last) = components.last() else {
276 continue;
277 };
278 by_last_component
279 .entry(last.clone())
280 .or_default()
281 .push(ProjectRootModuleName { file, components });
282 }
283 let mut path_extensions: HashMap<PathBuf, HashSet<PathBuf>> = HashMap::default();
284 for marker in package_markers {
285 let Some(directory) = marker.rel_path().parent() else {
286 continue;
287 };
288 if !python_package_holds_a_repeat_of_itself(directory, &package_dirs) {
289 continue;
290 }
291 let Some(source) = read_source(&marker) else {
292 continue;
293 };
294 let extended = python_appended_path_dirs(&source, marker.rel_path());
295 if !extended.is_empty() {
296 path_extensions
297 .entry(directory.to_path_buf())
298 .or_default()
299 .extend(extended);
300 }
301 }
302 Self {
303 by_last_component,
304 package_dirs,
305 path_extensions,
306 }
307 }
308
309 pub fn suffix_candidates(&self, module_fq: &str) -> Vec<PythonSourceRootModule> {
321 if module_fq.starts_with('.') {
325 return Vec::new();
326 }
327 let components = parse_symbol_path(Language::Python, module_fq);
328 let Some(last) = components.last().filter(|_| components.len() > 1) else {
329 return Vec::new();
330 };
331 let mut modules = Vec::new();
332 for candidate in self
333 .by_last_component
334 .get(last.as_str())
335 .into_iter()
336 .flatten()
337 {
338 let Some(dropped) = candidate
341 .components
342 .len()
343 .checked_sub(components.len())
344 .filter(|dropped| *dropped > 0)
345 else {
346 continue;
347 };
348 if candidate.components[dropped..] != components[..] {
349 continue;
350 }
351 let source_root: PathBuf = candidate
354 .file
355 .rel_path()
356 .components()
357 .take(dropped)
358 .collect();
359 debug_assert!(
360 !source_root.as_os_str().is_empty(),
361 "dropping a component of {:?} leaves a directory",
362 candidate.file.rel_path()
363 );
364 modules.push(PythonSourceRootModule {
365 file: candidate.file.clone(),
366 source_root,
367 });
368 }
369 modules
370 }
371
372 pub fn establishes_source_root(&self, candidate: &PythonSourceRootModule) -> bool {
418 if let Some(declared) = python_declared_import_root(&candidate.file) {
419 if declared == candidate.source_root {
420 return true;
421 }
422 let Some(extension) = python_repeated_root_extension_dir(
423 &candidate.file,
424 &candidate.source_root,
425 &declared,
426 ) else {
427 return false;
428 };
429 return self.package_chain_reaches_module(candidate)
430 && self
431 .path_extensions
432 .get(&candidate.source_root)
433 .is_some_and(|extended| extended.contains(&extension));
434 }
435 self.package_chain_reaches_module(candidate)
436 }
437
438 fn package_chain_reaches_module(&self, candidate: &PythonSourceRootModule) -> bool {
447 debug_assert!(
448 candidate
449 .file
450 .rel_path()
451 .starts_with(&candidate.source_root),
452 "a candidate's implied root is a prefix of its own path"
453 );
454 let mut directory = candidate.file.rel_path().parent();
455 while let Some(current) = directory {
456 if !self.package_dirs.contains(current) {
457 return false;
458 }
459 if current == candidate.source_root {
460 return true;
461 }
462 directory = current.parent();
463 }
464 false
465 }
466}
467
468fn python_package_holds_a_repeat_of_itself(
476 directory: &Path,
477 package_dirs: &HashSet<PathBuf>,
478) -> bool {
479 directory
480 .file_name()
481 .is_some_and(|name| package_dirs.contains(&directory.join(name)))
482}
483
484fn python_appended_path_dirs(source: &str, marker_rel_path: &Path) -> Vec<PathBuf> {
499 let Some(tree) = parse_python_tree(source) else {
500 return Vec::new();
501 };
502 let mut bound: HashMap<String, PathBuf> = HashMap::default();
503 bound.insert("__file__".to_owned(), marker_rel_path.to_path_buf());
504 let mut appended: Vec<PathBuf> = Vec::new();
505 walk_named_tree_preorder(tree.root_node(), true, |node| {
506 match node.kind() {
507 "function_definition" | "class_definition" | "lambda" => {
511 return WalkControl::SkipChildren;
512 }
513 "assignment" => {
514 if python_names_path(node.child_by_field_name("left"), source) {
519 appended.clear();
520 }
521 python_bind_assignment(node, source, &mut bound);
522 }
523 "call" => {
524 let directories = python_path_appended_dirs_in_call(node, source, &bound);
525 if directories.is_empty()
526 || !python_extension_statement_runs(node, source, &bound, &directories)
527 {
528 return WalkControl::Continue;
529 }
530 for directory in directories {
531 if directory.is_relative()
532 && !directory.as_os_str().is_empty()
533 && !appended.contains(&directory)
534 {
535 appended.push(directory);
536 }
537 }
538 }
539 _ => {}
540 }
541 WalkControl::Continue
542 });
543 appended
544}
545
546fn python_names_path(node: Option<Node<'_>>, source: &str) -> bool {
548 node.is_some_and(|node| {
549 node.kind() == "identifier" && node.utf8_text(source.as_bytes()).ok() == Some("__path__")
550 })
551}
552
553fn python_extension_statement_runs(
568 call: Node<'_>,
569 source: &str,
570 bound: &HashMap<String, PathBuf>,
571 directories: &[PathBuf],
572) -> bool {
573 let Some(statement) = call.parent() else {
574 return false;
575 };
576 if statement.kind() != "expression_statement" {
577 return false;
578 }
579 let Some(container) = statement.parent() else {
580 return false;
581 };
582 if container.kind() == "module" {
583 return true;
584 }
585 let guarded = if container.kind() == "block" {
590 container.parent().filter(|parent| {
591 parent.kind() == "if_statement"
592 && parent.child_by_field_name("consequence") == Some(container)
593 })
594 } else if container.kind() == "if_statement"
595 && container.child_by_field_name("consequence") == Some(statement)
596 {
597 Some(container)
598 } else {
599 None
600 };
601 let Some(guarded) = guarded else {
602 return false;
603 };
604 if guarded.parent().map(|parent| parent.kind()) != Some("module") {
609 return false;
610 }
611 let Some(condition) = guarded.child_by_field_name("condition") else {
612 return false;
613 };
614 python_condition_justifies_path_extension(condition, source, bound, directories)
615}
616
617fn python_condition_justifies_path_extension(
622 condition: Node<'_>,
623 source: &str,
624 bound: &HashMap<String, PathBuf>,
625 directories: &[PathBuf],
626) -> bool {
627 if condition.kind() != "comparison_operator" {
628 return false;
629 }
630 let mut cursor = condition.walk();
631 let operators = condition
632 .children_by_field_name("operators", &mut cursor)
633 .map(|operator| operator.kind())
634 .collect::<Vec<_>>();
635 if operators.as_slice() != ["not in"] {
636 return false;
637 }
638 let mut cursor = condition.walk();
641 let operands: Vec<Node<'_>> = condition.named_children(&mut cursor).collect();
642 let [left, right] = operands.as_slice() else {
643 return false;
644 };
645 if !python_names_path(Some(*right), source) {
646 return false;
647 }
648 python_path_expression_dir(*left, source, bound)
649 .is_some_and(|directory| directories.contains(&directory))
650}
651
652fn python_bind_assignment(
657 assignment: Node<'_>,
658 source: &str,
659 bound: &mut HashMap<String, PathBuf>,
660) {
661 let mut targets: Vec<Node<'_>> = Vec::new();
662 let mut current = assignment;
663 loop {
664 let (Some(left), Some(right)) = (
665 current.child_by_field_name("left"),
666 current.child_by_field_name("right"),
667 ) else {
668 return;
669 };
670 targets.push(left);
671 if right.kind() != "assignment" {
672 let resolved = python_path_expression_dir(right, source, bound);
673 for target in targets {
674 if target.kind() != "identifier" {
675 continue;
676 }
677 let Ok(name) = target.utf8_text(source.as_bytes()) else {
678 continue;
679 };
680 match &resolved {
681 Some(directory) => {
682 bound.insert(name.to_owned(), directory.clone());
683 }
684 None => {
685 bound.remove(name);
686 }
687 }
688 }
689 return;
690 }
691 current = right;
692 }
693}
694
695fn python_path_appended_dirs_in_call(
702 call: Node<'_>,
703 source: &str,
704 bound: &HashMap<String, PathBuf>,
705) -> Vec<PathBuf> {
706 let Some(function) = call.child_by_field_name("function") else {
707 return Vec::new();
708 };
709 let Some(path) = python_static_attribute_path(function) else {
710 return Vec::new();
711 };
712 let [object, method] = path.as_slice() else {
713 return Vec::new();
714 };
715 if object.utf8_text(source.as_bytes()).ok() != Some("__path__") {
716 return Vec::new();
717 }
718 let Ok(method) = method.utf8_text(source.as_bytes()) else {
719 return Vec::new();
720 };
721 let Some(arguments) = call.child_by_field_name("arguments") else {
722 return Vec::new();
723 };
724 let mut cursor = arguments.walk();
725 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
726 let named: Vec<Node<'_>> = match method {
727 "append" => arguments.first().copied().into_iter().collect(),
728 "insert" => arguments.get(1).copied().into_iter().collect(),
729 "extend" => match arguments.first() {
730 Some(list) if list.kind() == "list" => {
731 let mut cursor = list.walk();
732 list.named_children(&mut cursor).collect()
733 }
734 Some(iterator) => vec![*iterator],
735 None => Vec::new(),
736 },
737 _ => Vec::new(),
738 };
739 named
740 .into_iter()
741 .filter_map(|node| python_path_expression_dir(node, source, bound))
742 .collect()
743}
744
745fn python_path_expression_dir(
749 node: Node<'_>,
750 source: &str,
751 bound: &HashMap<String, PathBuf>,
752) -> Option<PathBuf> {
753 if let Some(literal) = python_plain_string_literal(node, source) {
754 let directory = PathBuf::from(literal);
755 return (!directory.as_os_str().is_empty()).then_some(directory);
756 }
757 match node.kind() {
758 "identifier" => bound.get(node.utf8_text(source.as_bytes()).ok()?).cloned(),
759 "call" => python_path_call_dir(node, source, bound),
760 _ => None,
761 }
762}
763
764fn python_path_call_dir(
767 call: Node<'_>,
768 source: &str,
769 bound: &HashMap<String, PathBuf>,
770) -> Option<PathBuf> {
771 let path = python_static_attribute_path(call.child_by_field_name("function")?)?;
772 let names: Vec<&str> = path
773 .iter()
774 .map(|node| node.utf8_text(source.as_bytes()).ok())
775 .collect::<Option<Vec<_>>>()?;
776 let arguments = call.child_by_field_name("arguments")?;
777 let mut cursor = arguments.walk();
778 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
779 match names.as_slice() {
780 ["os", "path", "dirname"] => {
781 let directory = python_path_expression_dir(*arguments.first()?, source, bound)?;
782 Some(directory.parent()?.to_path_buf())
783 }
784 ["os", "path", "join"] => {
785 let mut parts = arguments.iter();
786 let mut joined = python_path_expression_dir(*parts.next()?, source, bound)?;
787 for part in parts {
788 joined.push(python_path_expression_dir(*part, source, bound)?);
789 }
790 Some(joined)
791 }
792 _ => None,
793 }
794}
795
796pub fn resolve_module_code_unit(
804 python: &dyn PythonSource,
805 importer: Option<&ProjectFile>,
806 module_fq: &str,
807) -> Option<CodeUnit> {
808 let indexed = match python.path_module_fqn(module_fq) {
809 Some(units) => indexed_module_for_importer(importer, units),
812 None => indexed_module_for_importer(importer, python.definition_fqn(module_fq)),
813 };
814 match indexed {
815 IndexedModuleLookup::Unique(module) => Some(module),
816 IndexedModuleLookup::Ambiguous => None,
820 IndexedModuleLookup::Missing => {
821 let importer = importer?;
822 module_under_source_root(&python.module_spellings(), importer, module_fq)
823 }
824 }
825}
826
827enum IndexedModuleLookup {
836 Unique(CodeUnit),
837 Ambiguous,
838 Missing,
839}
840
841fn indexed_module_for_importer(
842 importer: Option<&ProjectFile>,
843 units: Vec<CodeUnit>,
844) -> IndexedModuleLookup {
845 let mut modules = units
846 .into_iter()
847 .filter(CodeUnit::is_module)
848 .collect::<Vec<_>>();
849 if let Some(importer) = importer {
850 retain_modules_for_importer(importer, &mut modules, CodeUnit::source);
851 }
852 match modules.len() {
853 0 => IndexedModuleLookup::Missing,
854 1 => IndexedModuleLookup::Unique(modules.remove(0)),
855 _ => IndexedModuleLookup::Ambiguous,
856 }
857}
858
859pub fn python_module_file_under_source_root(
888 spellings: &PythonModuleSpellings,
889 importer: &ProjectFile,
890 module_fq: &str,
891) -> Option<ProjectFile> {
892 let mut candidates = spellings.suffix_candidates(module_fq);
893 retain_nearest_owning_root(importer, &mut candidates, |candidate| {
894 candidate.source_root.clone()
895 });
896 candidates.retain(|candidate| spellings.establishes_source_root(candidate));
897 let [candidate] = candidates.as_slice() else {
898 return None;
899 };
900 debug_assert_eq!(
901 importer.root(),
902 candidate.file.root(),
903 "one workspace's Python files share one project root"
904 );
905 Some(candidate.file.clone())
906}
907
908fn module_under_source_root(
913 spellings: &PythonModuleSpellings,
914 importer: &ProjectFile,
915 module_fq: &str,
916) -> Option<CodeUnit> {
917 let file = python_module_file_under_source_root(spellings, importer, module_fq)?;
918 module_code_unit(&file, &python_module_name(&file))
919}
920
921pub fn resolve_module_code_units_batch(
927 python: &dyn PythonSource,
928 importer: Option<&ProjectFile>,
929 module_fqs: &[String],
930) -> Vec<Option<CodeUnit>> {
931 let path_results = python.path_module_fqns_batch(module_fqs);
932 let mut lookups: Vec<IndexedModuleLookup> = Vec::with_capacity(module_fqs.len());
933 let mut needs_definition_fallback = Vec::new();
934 for (i, units) in path_results.into_iter().enumerate() {
935 match units {
936 Some(units) => lookups.push(indexed_module_for_importer(importer, units)),
937 None => {
938 needs_definition_fallback.push(i);
939 lookups.push(IndexedModuleLookup::Missing);
941 }
942 }
943 }
944 for i in needs_definition_fallback {
945 lookups[i] = indexed_module_for_importer(importer, python.definition_fqn(&module_fqs[i]));
946 }
947 let mut results: Vec<Option<CodeUnit>> = vec![None; module_fqs.len()];
948 let mut needs_source_root = Vec::new();
949 for (i, lookup) in lookups.into_iter().enumerate() {
950 match lookup {
951 IndexedModuleLookup::Unique(module) => results[i] = Some(module),
952 IndexedModuleLookup::Ambiguous => {}
953 IndexedModuleLookup::Missing => needs_source_root.push(i),
954 }
955 }
956 if let Some(importer) = importer
957 && !needs_source_root.is_empty()
958 {
959 let spellings = python.module_spellings();
962 for i in needs_source_root {
963 results[i] = module_under_source_root(&spellings, importer, &module_fqs[i]);
964 }
965 }
966 results
967}
968
969pub fn compute_export_index_of(
970 python: &dyn PythonSource,
971 token: QueryToken<'_>,
972 file: &ProjectFile,
973) -> ExportIndex {
974 let mut index = ExportIndex::empty();
975 let mut events = Vec::new();
976 let declarations = python.top_level_declarations(file);
977 collect_local_export_events(
978 declarations.iter(),
979 |code_unit| {
980 python
981 .ranges(code_unit)
982 .iter()
983 .map(|range| range.start_byte)
984 .min()
985 .unwrap_or(usize::MAX)
986 },
987 &mut events,
988 );
989
990 if let Ok(source) = file.read_to_string()
991 && let Some(tree) = parse_python_tree(&source)
992 {
993 collect_reexport_events(
994 python,
995 file,
996 tree.root_node(),
997 &source,
998 &mut events,
999 &mut index,
1000 );
1001 } else {
1002 let imports = python.import_info_of(token, file);
1003 collect_reexport_events_from_imports(python, file, &imports, &mut events, &mut index);
1004 }
1005
1006 finish_export_index(events, index)
1007}
1008
1009pub fn export_index_from_file_facts(
1010 python: &dyn PythonSource,
1011 file: &ProjectFile,
1012 facts: &dyn IndexedFileFacts,
1013 module_name: &str,
1014 binder: &ImportBinder,
1015) -> ExportIndex {
1016 let mut index = ExportIndex::empty();
1017 let mut events = Vec::new();
1018 let mut local_names = collect_local_export_events(
1019 facts.top_level_declarations().iter(),
1020 |code_unit| {
1021 facts
1022 .declaration_ranges(code_unit)
1023 .into_iter()
1024 .flatten()
1025 .map(|range| range.start_byte)
1026 .min()
1027 .unwrap_or(usize::MAX)
1028 },
1029 &mut events,
1030 );
1031
1032 if !facts
1033 .top_level_declarations()
1034 .iter()
1035 .any(CodeUnit::is_module)
1036 && let Some(identifier) = module_name.rsplit('.').next()
1037 && !identifier.is_empty()
1038 && !identifier.starts_with('_')
1039 {
1040 local_names.insert(identifier.to_string());
1041 events.push((
1042 0,
1043 identifier.to_string(),
1044 ExportEntry::Local {
1045 local_name: identifier.to_string(),
1046 },
1047 ));
1048 }
1049
1050 if import_order_requires_source(binder, &local_names)
1051 && let Ok(source) = file.read_to_string()
1052 && let Some(tree) = parse_python_tree(&source)
1053 {
1054 collect_reexport_events(
1055 python,
1056 file,
1057 tree.root_node(),
1058 &source,
1059 &mut events,
1060 &mut index,
1061 );
1062 } else {
1063 collect_reexport_events_from_imports(
1064 python,
1065 file,
1066 facts.imports(),
1067 &mut events,
1068 &mut index,
1069 );
1070 }
1071
1072 finish_export_index(events, index)
1073}
1074
1075fn collect_local_export_events<'a>(
1076 declarations: impl IntoIterator<Item = &'a CodeUnit>,
1077 mut start_byte: impl FnMut(&CodeUnit) -> usize,
1078 events: &mut Vec<(usize, String, ExportEntry)>,
1079) -> HashSet<String> {
1080 let mut local_names = HashSet::default();
1081 for code_unit in declarations {
1082 let identifier = code_unit.identifier().trim();
1083 if identifier.is_empty() {
1084 continue;
1085 }
1086 local_names.insert(identifier.to_string());
1087 events.push((
1088 start_byte(code_unit),
1089 identifier.to_string(),
1090 ExportEntry::Local {
1091 local_name: identifier.to_string(),
1092 },
1093 ));
1094 }
1095 local_names
1096}
1097
1098fn finish_export_index(
1099 mut events: Vec<(usize, String, ExportEntry)>,
1100 mut index: ExportIndex,
1101) -> ExportIndex {
1102 events.sort_by_key(|(start_byte, _, _)| *start_byte);
1103 for (_, exported_name, entry) in events {
1104 index.exports_by_name.insert(exported_name, entry);
1105 }
1106 index
1107}
1108
1109fn collect_reexport_events(
1110 python: &dyn PythonSource,
1111 file: &ProjectFile,
1112 root: tree_sitter::Node<'_>,
1113 source: &str,
1114 events: &mut Vec<(usize, String, ExportEntry)>,
1115 index: &mut ExportIndex,
1116) {
1117 let mut stack = vec![root];
1122 while let Some(node) = stack.pop() {
1123 let mut cursor = node.walk();
1124 for child in node.named_children(&mut cursor) {
1125 match child.kind() {
1126 "import_from_statement" => {
1127 for info in python_import_infos_from_node(child, source) {
1128 record_single_reexport_event(python, file, &info, events, index);
1129 }
1130 }
1131 "function_definition" | "class_definition" => {}
1132 _ => stack.push(child),
1133 }
1134 }
1135 }
1136}
1137
1138fn collect_reexport_events_from_imports(
1139 python: &dyn PythonSource,
1140 file: &ProjectFile,
1141 imports: &[ImportInfo],
1142 events: &mut Vec<(usize, String, ExportEntry)>,
1143 index: &mut ExportIndex,
1144) {
1145 for import in imports {
1146 record_single_reexport_event(python, file, import, events, index);
1147 }
1148}
1149
1150fn record_single_reexport_event(
1151 python: &dyn PythonSource,
1152 file: &ProjectFile,
1153 import: &ImportInfo,
1154 events: &mut Vec<(usize, String, ExportEntry)>,
1155 index: &mut ExportIndex,
1156) {
1157 let Some(PythonImportDetails::FromImport {
1158 module,
1159 name,
1160 alias,
1161 wildcard,
1162 }) = python_import_details(import)
1163 else {
1164 return;
1165 };
1166 let start_byte = import
1167 .path
1168 .as_ref()
1169 .map(|path| path.declaration_start_byte)
1170 .unwrap_or(usize::MAX);
1171 let resolved_module = if module.starts_with('.') {
1172 resolve_python_relative_module(file, &module)
1173 } else {
1174 Some(module.clone())
1175 };
1176 let Some(resolved_module) = resolved_module else {
1177 return;
1178 };
1179
1180 if wildcard {
1181 index.reexport_stars.push(ReexportStar {
1182 module_specifier: resolved_module,
1183 });
1184 return;
1185 }
1186 let exported_name = alias.unwrap_or(name.clone());
1187 let module_candidate = format!("{resolved_module}.{name}");
1193 if resolve_module_code_unit(python, Some(file), &module_candidate).is_some() {
1194 events.push((
1195 start_byte,
1196 exported_name,
1197 ExportEntry::ReexportedModule {
1198 module_specifier: module_candidate,
1199 },
1200 ));
1201 return;
1202 }
1203 events.push((
1204 start_byte,
1205 exported_name,
1206 ExportEntry::ReexportedNamed {
1207 module_specifier: resolved_module,
1208 imported_name: name,
1209 },
1210 ));
1211}
1212
1213pub fn import_binder_from_imports(
1214 python: &dyn PythonSource,
1215 file: &ProjectFile,
1216 imports: &[ImportInfo],
1217) -> ImportBinder {
1218 let mut binder = ImportBinder::empty();
1219
1220 for (local_name, binding) in import_bindings_from_imports(python, file, imports) {
1221 binder.bindings.insert(local_name, binding);
1222 }
1223
1224 binder
1225}
1226
1227pub fn import_bindings_from_imports(
1232 python: &dyn PythonSource,
1233 file: &ProjectFile,
1234 imports: &[ImportInfo],
1235) -> Vec<(String, ImportBinding)> {
1236 let mut bindings = Vec::new();
1237
1238 for import in imports {
1239 let Some(details) = python_import_details(import) else {
1240 continue;
1241 };
1242 match details {
1243 PythonImportDetails::Import { module, alias } => {
1244 let local_name = python_namespace_binding_name(import, alias.as_deref(), &module);
1245 let module_specifier =
1246 python_namespace_binding_module(import, alias.as_deref(), &module);
1247 bindings.push((
1248 local_name,
1249 ImportBinding {
1250 module_specifier,
1251 namespace_imported_module: Some(module),
1252 kind: ImportKind::Namespace,
1253 imported_name: None,
1254 },
1255 ));
1256 }
1257 PythonImportDetails::FromImport {
1258 module,
1259 name,
1260 wildcard,
1261 ..
1262 } => {
1263 let resolved_module = if module.starts_with('.') {
1264 resolve_python_relative_module(file, &module)
1265 } else {
1266 Some(module.clone())
1267 };
1268 let Some(resolved_module) = resolved_module else {
1269 continue;
1270 };
1271 if wildcard {
1272 bindings.extend(
1277 public_declarations_in_module(python, file, &resolved_module)
1278 .into_iter()
1279 .map(|declaration| {
1280 let name = declaration.identifier().to_string();
1281 (
1282 name.clone(),
1283 ImportBinding {
1284 module_specifier: resolved_module.clone(),
1285 namespace_imported_module: None,
1286 kind: ImportKind::Named,
1287 imported_name: Some(name),
1288 },
1289 )
1290 }),
1291 );
1292 continue;
1293 }
1294 let local_name = import
1299 .local_name()
1300 .map(str::to_string)
1301 .unwrap_or_else(|| name.clone());
1302 let module_candidate = format!("{resolved_module}.{name}");
1303 if resolve_module_code_unit(python, Some(file), &module_candidate).is_some() {
1304 bindings.push((
1305 local_name,
1306 ImportBinding {
1307 module_specifier: module_candidate,
1308 namespace_imported_module: None,
1309 kind: ImportKind::Namespace,
1310 imported_name: None,
1311 },
1312 ));
1313 continue;
1314 }
1315 bindings.push((
1316 local_name,
1317 ImportBinding {
1318 module_specifier: resolved_module,
1319 namespace_imported_module: None,
1320 kind: ImportKind::Named,
1321 imported_name: Some(name),
1322 },
1323 ));
1324 }
1325 }
1326 }
1327
1328 bindings
1329}
1330
1331pub fn public_declarations_in_module(
1332 python: &dyn PythonSource,
1333 importer: &ProjectFile,
1334 module_fq: &str,
1335) -> Vec<CodeUnit> {
1336 let Some(module_code_unit) = resolve_module_code_unit(python, Some(importer), module_fq) else {
1337 return Vec::new();
1338 };
1339 python
1340 .direct_children(&module_code_unit)
1341 .into_iter()
1342 .filter(|code_unit| !code_unit.identifier().starts_with('_'))
1343 .collect()
1344}
1345
1346pub fn resolve_base_class(
1347 python: &dyn PythonSource,
1348 token: QueryToken<'_>,
1349 code_unit: &CodeUnit,
1350 raw: &str,
1351) -> Option<CodeUnit> {
1352 let trimmed = raw.trim();
1353 if trimmed.is_empty() {
1354 return None;
1355 }
1356
1357 let binder = python.import_binder_of(code_unit.source());
1358 if let Some((head, tail)) = trimmed.split_once('.') {
1359 if let Some(binding) = binder.bindings.get(head)
1360 && binding.kind == ImportKind::Namespace
1361 {
1362 let fq_name = format!("{}.{}", binding.module_specifier, tail);
1363 return python.definitions(&fq_name).next();
1364 }
1365 return python.definitions(trimmed).next();
1366 }
1367
1368 if let Some(binding) = binder.bindings.get(trimmed) {
1369 match binding.kind {
1370 ImportKind::Namespace => {
1371 return resolve_module_code_unit(
1372 python,
1373 Some(code_unit.source()),
1374 &binding.module_specifier,
1375 );
1376 }
1377 ImportKind::Named => {
1378 let imported_name = binding.imported_name.as_ref()?;
1379 let fqn = format!("{}.{}", binding.module_specifier, imported_name);
1380 return resolve_exported_fqn(python, Some(code_unit.source()), &fqn)
1381 .into_iter()
1382 .next()
1383 .or_else(|| python.definitions(&fqn).next());
1384 }
1385 _ => {}
1386 }
1387 }
1388
1389 if python
1390 .import_info_of(token, code_unit.source())
1391 .iter()
1392 .any(|import| import.is_wildcard)
1393 && let Some(imported) =
1394 resolve_import_bindings(python, token, code_unit.source()).get(trimmed)
1395 {
1396 return Some(imported.clone());
1397 }
1398
1399 let local_fq_name = format!("{}.{}", code_unit.package_name(), trimmed);
1400 python
1401 .definitions(&local_fq_name)
1402 .next()
1403 .or_else(|| python.definitions(trimmed).next())
1404}
1405
1406pub fn render_skeleton_recursive(
1407 index: &dyn CodeUnitIndex,
1408 code_unit: &CodeUnit,
1409 indent: &str,
1410 header_only: bool,
1411 out: &mut String,
1412) {
1413 if let Some(signature) = python_signature(index, code_unit, header_only) {
1414 for line in signature.lines() {
1415 out.push_str(indent);
1416 out.push_str(line);
1417 out.push('\n');
1418 }
1419 }
1420
1421 let all_children = index.direct_children(code_unit);
1422 let field_children: Vec<_> = all_children
1423 .iter()
1424 .filter(|child| child.is_field())
1425 .cloned()
1426 .collect();
1427 let children = if header_only {
1428 field_children.clone()
1429 } else {
1430 all_children.clone()
1431 };
1432 if !children.is_empty() || code_unit.is_class() || code_unit.is_module() {
1433 let child_indent = format!("{indent} ");
1434 for child in children {
1435 render_skeleton_recursive(index, &child, &child_indent, header_only, out);
1436 }
1437 if header_only && all_children.len() > field_children.len() {
1438 out.push_str(&child_indent);
1439 out.push_str("[...]\n");
1440 }
1441 }
1442}
1443
1444fn python_signature(
1445 index: &dyn CodeUnitIndex,
1446 code_unit: &CodeUnit,
1447 _header_only: bool,
1448) -> Option<String> {
1449 if code_unit.is_module() {
1450 return None;
1451 }
1452
1453 let source = index.get_source(code_unit, false)?;
1454 let lines: Vec<_> = source
1455 .lines()
1456 .map(str::trim_end)
1457 .filter(|line| !line.trim().is_empty())
1458 .collect();
1459 if lines.is_empty() {
1460 return None;
1461 }
1462
1463 let mut decorators = Vec::new();
1464 let mut header = None;
1465 for line in lines {
1466 let trimmed = line.trim_start();
1467 if trimmed.starts_with('@') {
1468 decorators.push(trimmed.to_string());
1469 continue;
1470 }
1471 header = Some(trimmed.to_string());
1472 break;
1473 }
1474 let mut rendered = String::new();
1475 for decorator in decorators {
1476 rendered.push_str(&decorator);
1477 rendered.push('\n');
1478 }
1479
1480 let header = header?;
1481 match code_unit.kind() {
1482 CodeUnitType::Class => rendered.push_str(&header),
1483 CodeUnitType::Function => {
1484 rendered.push_str(header.trim_end_matches(':'));
1485 rendered.push_str(": ...");
1486 }
1487 CodeUnitType::Field | CodeUnitType::Macro => rendered.push_str(header.as_str()),
1488 CodeUnitType::Module | CodeUnitType::FileScope => return None,
1489 }
1490 Some(rendered)
1491}
1492
1493fn import_order_requires_source(binder: &ImportBinder, local_names: &HashSet<String>) -> bool {
1494 binder
1495 .bindings
1496 .keys()
1497 .any(|bound_name| local_names.contains(bound_name))
1498}