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 record_memo_inputs(&self) -> [u8; 32];
58
59 fn path_module_fqn(&self, module_fq: &str) -> Option<Vec<CodeUnit>>;
60
61 fn path_module_fqns_batch(&self, module_fqs: &[String]) -> Vec<Option<Vec<CodeUnit>>>;
64
65 fn definition_fqn(&self, fqn: &str) -> Vec<CodeUnit>;
66
67 fn module_spellings(&self) -> Arc<PythonModuleSpellings>;
74
75 fn import_binder_of(&self, file: &ProjectFile) -> Arc<ImportBinder>;
80
81 fn export_index_of(&self, file: &ProjectFile) -> Arc<ExportIndex>;
82
83 fn prepared_syntax(
88 &self,
89 token: QueryToken<'_>,
90 file: &ProjectFile,
91 ) -> Option<Arc<PreparedSyntaxTree>>;
92
93 fn visit_file_facts(
96 &self,
97 files: &[ProjectFile],
98 keep_going: &dyn Fn() -> bool,
99 visit: &mut dyn FnMut(&ProjectFile, Option<&dyn IndexedFileFacts>),
100 );
101}
102
103pub trait PythonUsageSource: PythonSource {
106 fn usage_index(&self) -> Arc<PythonUsageIndex>;
107 fn module_binding_timeline(
108 &self,
109 file: &ProjectFile,
110 build: &mut dyn FnMut() -> crate::usage_index::ModuleBindingTimeline,
111 ) -> Arc<crate::usage_index::ModuleBindingTimeline>;
112 fn scope_facts(
113 &self,
114 file: &ProjectFile,
115 build: &mut dyn FnMut() -> Option<crate::usage_index::PythonScopeFacts>,
116 ) -> Option<Arc<crate::usage_index::PythonScopeFacts>>;
117
118 fn usage_index_with_cancellation(
120 &self,
121 cancellation: &brokk_bifrost_core::CancellationToken,
122 ) -> Option<Arc<PythonUsageIndex>>;
123}
124
125pub fn retain_modules_for_importer<T>(
144 project: &dyn Project,
145 importer: &ProjectFile,
146 candidates: &mut Vec<T>,
147 source_of: impl Fn(&T) -> &ProjectFile,
148) {
149 retain_nearest_owning_root(importer, candidates, |candidate| {
150 let source = source_of(candidate);
151 debug_assert_eq!(
152 importer.root(),
153 source.root(),
154 "one workspace's Python files share one project root"
155 );
156 python_import_root(project, source)
157 });
158}
159
160fn retain_nearest_owning_root<T>(
169 importer: &ProjectFile,
170 candidates: &mut Vec<T>,
171 root_of: impl Fn(&T) -> PathBuf,
172) {
173 if candidates.len() < 2 {
174 return;
175 }
176 let owning_depth = |candidate: &T| {
177 let root = root_of(candidate);
178 importer
179 .rel_path()
180 .starts_with(&root)
181 .then(|| root.components().count())
182 };
183 let depths = candidates.iter().map(owning_depth).collect::<Vec<_>>();
184 let Some(nearest) = depths.iter().copied().flatten().max() else {
185 return;
186 };
187 *candidates = std::mem::take(candidates)
188 .into_iter()
189 .zip(depths)
190 .filter(|(_, depth)| *depth == Some(nearest))
191 .map(|(candidate, _)| candidate)
192 .collect();
193}
194
195#[derive(Clone, Debug)]
197struct ProjectRootModuleName {
198 file: ProjectFile,
199 components: Vec<String>,
200}
201
202#[derive(Clone, Debug)]
211pub struct PythonSourceRootModule {
212 pub file: ProjectFile,
213 pub source_root: PathBuf,
217}
218
219#[derive(Default)]
241pub struct PythonModuleSpellings {
242 project: Option<Arc<dyn Project>>,
243 by_last_component: HashMap<String, Vec<ProjectRootModuleName>>,
244 package_dirs: HashSet<PathBuf>,
247 path_extensions: HashMap<PathBuf, HashSet<PathBuf>>,
259}
260
261impl std::fmt::Debug for PythonModuleSpellings {
262 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263 f.debug_struct("PythonModuleSpellings")
264 .field("by_last_component", &self.by_last_component)
265 .field("package_dirs", &self.package_dirs)
266 .field("path_extensions", &self.path_extensions)
267 .finish_non_exhaustive()
268 }
269}
270
271impl PythonModuleSpellings {
272 pub fn authority_rows(&self) -> Vec<String> {
275 let mut rows = Vec::new();
276 for (last, candidates) in &self.by_last_component {
277 for candidate in candidates {
278 let root = self
279 .project()
280 .and_then(|project| python_declared_import_root(project, &candidate.file))
281 .map(|path| {
282 brokk_bifrost_core::path_utils::workspace_relative_slash_path(&path)
283 });
284 rows.push(format!(
285 "module:{last:?}:{:?}:{:?}:{root:?}",
286 brokk_bifrost_core::path_utils::rel_path_string(&candidate.file),
287 candidate.components
288 ));
289 }
290 }
291 rows.extend(self.package_dirs.iter().map(|directory| {
292 format!(
293 "package:{:?}",
294 brokk_bifrost_core::path_utils::workspace_relative_slash_path(directory)
295 )
296 }));
297 for (directory, extensions) in &self.path_extensions {
298 for extension in extensions {
299 rows.push(format!(
300 "extension:{:?}:{:?}",
301 brokk_bifrost_core::path_utils::workspace_relative_slash_path(directory),
302 brokk_bifrost_core::path_utils::workspace_relative_slash_path(extension)
303 ));
304 }
305 }
306 rows.sort();
307 rows
308 }
309
310 pub(crate) fn project(&self) -> Option<&dyn Project> {
311 self.project.as_deref()
312 }
313
314 pub fn build(files: impl IntoIterator<Item = ProjectFile>, project: Arc<dyn Project>) -> Self {
320 let mut by_last_component: HashMap<String, Vec<ProjectRootModuleName>> = HashMap::default();
321 let mut package_dirs: HashSet<PathBuf> = HashSet::default();
322 let mut package_markers: Vec<ProjectFile> = Vec::new();
323 for file in files {
324 if file.rel_path().file_name().and_then(|name| name.to_str()) == Some("__init__.py")
325 && let Some(directory) = file.rel_path().parent()
326 {
327 package_dirs.insert(directory.to_path_buf());
328 package_markers.push(file.clone());
329 }
330 let components = python_project_module_components(&file);
331 let Some(last) = components.last() else {
332 continue;
333 };
334 by_last_component
335 .entry(last.clone())
336 .or_default()
337 .push(ProjectRootModuleName { file, components });
338 }
339 let mut path_extensions: HashMap<PathBuf, HashSet<PathBuf>> = HashMap::default();
340 for marker in package_markers {
341 let Some(directory) = marker.rel_path().parent() else {
342 continue;
343 };
344 if !python_package_holds_a_repeat_of_itself(directory, &package_dirs) {
345 continue;
346 }
347 let Ok(source) = project.read_source(&marker) else {
348 continue;
349 };
350 let extended = python_appended_path_dirs(&source, marker.rel_path());
351 if !extended.is_empty() {
352 path_extensions
353 .entry(directory.to_path_buf())
354 .or_default()
355 .extend(extended);
356 }
357 }
358 Self {
359 project: Some(project),
360 by_last_component,
361 package_dirs,
362 path_extensions,
363 }
364 }
365
366 pub fn suffix_candidates(&self, module_fq: &str) -> Vec<PythonSourceRootModule> {
378 if module_fq.starts_with('.') {
382 return Vec::new();
383 }
384 let components = parse_symbol_path(Language::Python, module_fq);
385 let Some(last) = components.last().filter(|_| components.len() > 1) else {
386 return Vec::new();
387 };
388 let mut modules = Vec::new();
389 for candidate in self
390 .by_last_component
391 .get(last.as_str())
392 .into_iter()
393 .flatten()
394 {
395 let Some(dropped) = candidate
398 .components
399 .len()
400 .checked_sub(components.len())
401 .filter(|dropped| *dropped > 0)
402 else {
403 continue;
404 };
405 if candidate.components[dropped..] != components[..] {
406 continue;
407 }
408 let source_root: PathBuf = candidate
411 .file
412 .rel_path()
413 .components()
414 .take(dropped)
415 .collect();
416 debug_assert!(
417 !source_root.as_os_str().is_empty(),
418 "dropping a component of {:?} leaves a directory",
419 candidate.file.rel_path()
420 );
421 modules.push(PythonSourceRootModule {
422 file: candidate.file.clone(),
423 source_root,
424 });
425 }
426 modules
427 }
428
429 pub fn establishes_source_root(&self, candidate: &PythonSourceRootModule) -> bool {
475 if let Some(declared) = python_declared_import_root(
476 self.project()
477 .expect("source-root candidates have a selected project"),
478 &candidate.file,
479 ) {
480 if declared == candidate.source_root {
481 return true;
482 }
483 let Some(extension) = python_repeated_root_extension_dir(
484 &candidate.file,
485 &candidate.source_root,
486 &declared,
487 ) else {
488 return false;
489 };
490 return self.package_chain_reaches_module(candidate)
491 && self
492 .path_extensions
493 .get(&candidate.source_root)
494 .is_some_and(|extended| extended.contains(&extension));
495 }
496 self.package_chain_reaches_module(candidate)
497 }
498
499 fn package_chain_reaches_module(&self, candidate: &PythonSourceRootModule) -> bool {
508 debug_assert!(
509 candidate
510 .file
511 .rel_path()
512 .starts_with(&candidate.source_root),
513 "a candidate's implied root is a prefix of its own path"
514 );
515 let mut directory = candidate.file.rel_path().parent();
516 while let Some(current) = directory {
517 if !self.package_dirs.contains(current) {
518 return false;
519 }
520 if current == candidate.source_root {
521 return true;
522 }
523 directory = current.parent();
524 }
525 false
526 }
527}
528
529fn python_package_holds_a_repeat_of_itself(
537 directory: &Path,
538 package_dirs: &HashSet<PathBuf>,
539) -> bool {
540 directory
541 .file_name()
542 .is_some_and(|name| package_dirs.contains(&directory.join(name)))
543}
544
545fn python_appended_path_dirs(source: &str, marker_rel_path: &Path) -> Vec<PathBuf> {
560 let Some(tree) = parse_python_tree(source) else {
561 return Vec::new();
562 };
563 let mut bound: HashMap<String, PathBuf> = HashMap::default();
564 bound.insert("__file__".to_owned(), marker_rel_path.to_path_buf());
565 let mut appended: Vec<PathBuf> = Vec::new();
566 walk_named_tree_preorder(tree.root_node(), true, |node| {
567 match node.kind() {
568 "function_definition" | "class_definition" | "lambda" => {
572 return WalkControl::SkipChildren;
573 }
574 "assignment" => {
575 if python_names_path(node.child_by_field_name("left"), source) {
580 appended.clear();
581 }
582 python_bind_assignment(node, source, &mut bound);
583 }
584 "call" => {
585 let directories = python_path_appended_dirs_in_call(node, source, &bound);
586 if directories.is_empty()
587 || !python_extension_statement_runs(node, source, &bound, &directories)
588 {
589 return WalkControl::Continue;
590 }
591 for directory in directories {
592 if directory.is_relative()
593 && !directory.as_os_str().is_empty()
594 && !appended.contains(&directory)
595 {
596 appended.push(directory);
597 }
598 }
599 }
600 _ => {}
601 }
602 WalkControl::Continue
603 });
604 appended
605}
606
607fn python_names_path(node: Option<Node<'_>>, source: &str) -> bool {
609 node.is_some_and(|node| {
610 node.kind() == "identifier" && node.utf8_text(source.as_bytes()).ok() == Some("__path__")
611 })
612}
613
614fn python_extension_statement_runs(
629 call: Node<'_>,
630 source: &str,
631 bound: &HashMap<String, PathBuf>,
632 directories: &[PathBuf],
633) -> bool {
634 let Some(statement) = call.parent() else {
635 return false;
636 };
637 if statement.kind() != "expression_statement" {
638 return false;
639 }
640 let Some(container) = statement.parent() else {
641 return false;
642 };
643 if container.kind() == "module" {
644 return true;
645 }
646 let guarded = if container.kind() == "block" {
651 container.parent().filter(|parent| {
652 parent.kind() == "if_statement"
653 && parent.child_by_field_name("consequence") == Some(container)
654 })
655 } else if container.kind() == "if_statement"
656 && container.child_by_field_name("consequence") == Some(statement)
657 {
658 Some(container)
659 } else {
660 None
661 };
662 let Some(guarded) = guarded else {
663 return false;
664 };
665 if guarded.parent().map(|parent| parent.kind()) != Some("module") {
670 return false;
671 }
672 let Some(condition) = guarded.child_by_field_name("condition") else {
673 return false;
674 };
675 python_condition_justifies_path_extension(condition, source, bound, directories)
676}
677
678fn python_condition_justifies_path_extension(
683 condition: Node<'_>,
684 source: &str,
685 bound: &HashMap<String, PathBuf>,
686 directories: &[PathBuf],
687) -> bool {
688 if condition.kind() != "comparison_operator" {
689 return false;
690 }
691 let mut cursor = condition.walk();
692 let operators = condition
693 .children_by_field_name("operators", &mut cursor)
694 .map(|operator| operator.kind())
695 .collect::<Vec<_>>();
696 if operators.as_slice() != ["not in"] {
697 return false;
698 }
699 let mut cursor = condition.walk();
702 let operands: Vec<Node<'_>> = condition.named_children(&mut cursor).collect();
703 let [left, right] = operands.as_slice() else {
704 return false;
705 };
706 if !python_names_path(Some(*right), source) {
707 return false;
708 }
709 python_path_expression_dir(*left, source, bound)
710 .is_some_and(|directory| directories.contains(&directory))
711}
712
713fn python_bind_assignment(
718 assignment: Node<'_>,
719 source: &str,
720 bound: &mut HashMap<String, PathBuf>,
721) {
722 let mut targets: Vec<Node<'_>> = Vec::new();
723 let mut current = assignment;
724 loop {
725 let (Some(left), Some(right)) = (
726 current.child_by_field_name("left"),
727 current.child_by_field_name("right"),
728 ) else {
729 return;
730 };
731 targets.push(left);
732 if right.kind() != "assignment" {
733 let resolved = python_path_expression_dir(right, source, bound);
734 for target in targets {
735 if target.kind() != "identifier" {
736 continue;
737 }
738 let Ok(name) = target.utf8_text(source.as_bytes()) else {
739 continue;
740 };
741 match &resolved {
742 Some(directory) => {
743 bound.insert(name.to_owned(), directory.clone());
744 }
745 None => {
746 bound.remove(name);
747 }
748 }
749 }
750 return;
751 }
752 current = right;
753 }
754}
755
756fn python_path_appended_dirs_in_call(
763 call: Node<'_>,
764 source: &str,
765 bound: &HashMap<String, PathBuf>,
766) -> Vec<PathBuf> {
767 let Some(function) = call.child_by_field_name("function") else {
768 return Vec::new();
769 };
770 let Some(path) = python_static_attribute_path(function) else {
771 return Vec::new();
772 };
773 let [object, method] = path.as_slice() else {
774 return Vec::new();
775 };
776 if object.utf8_text(source.as_bytes()).ok() != Some("__path__") {
777 return Vec::new();
778 }
779 let Ok(method) = method.utf8_text(source.as_bytes()) else {
780 return Vec::new();
781 };
782 let Some(arguments) = call.child_by_field_name("arguments") else {
783 return Vec::new();
784 };
785 let mut cursor = arguments.walk();
786 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
787 let named: Vec<Node<'_>> = match method {
788 "append" => arguments.first().copied().into_iter().collect(),
789 "insert" => arguments.get(1).copied().into_iter().collect(),
790 "extend" => match arguments.first() {
791 Some(list) if list.kind() == "list" => {
792 let mut cursor = list.walk();
793 list.named_children(&mut cursor).collect()
794 }
795 Some(iterator) => vec![*iterator],
796 None => Vec::new(),
797 },
798 _ => Vec::new(),
799 };
800 named
801 .into_iter()
802 .filter_map(|node| python_path_expression_dir(node, source, bound))
803 .collect()
804}
805
806fn python_path_expression_dir(
810 node: Node<'_>,
811 source: &str,
812 bound: &HashMap<String, PathBuf>,
813) -> Option<PathBuf> {
814 if let Some(literal) = python_plain_string_literal(node, source) {
815 let directory = PathBuf::from(literal);
816 return (!directory.as_os_str().is_empty()).then_some(directory);
817 }
818 match node.kind() {
819 "identifier" => bound.get(node.utf8_text(source.as_bytes()).ok()?).cloned(),
820 "call" => python_path_call_dir(node, source, bound),
821 _ => None,
822 }
823}
824
825fn python_path_call_dir(
828 call: Node<'_>,
829 source: &str,
830 bound: &HashMap<String, PathBuf>,
831) -> Option<PathBuf> {
832 let path = python_static_attribute_path(call.child_by_field_name("function")?)?;
833 let names: Vec<&str> = path
834 .iter()
835 .map(|node| node.utf8_text(source.as_bytes()).ok())
836 .collect::<Option<Vec<_>>>()?;
837 let arguments = call.child_by_field_name("arguments")?;
838 let mut cursor = arguments.walk();
839 let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
840 match names.as_slice() {
841 ["os", "path", "dirname"] => {
842 let directory = python_path_expression_dir(*arguments.first()?, source, bound)?;
843 Some(directory.parent()?.to_path_buf())
844 }
845 ["os", "path", "join"] => {
846 let mut parts = arguments.iter();
847 let mut joined = python_path_expression_dir(*parts.next()?, source, bound)?;
848 for part in parts {
849 joined.push(python_path_expression_dir(*part, source, bound)?);
850 }
851 Some(joined)
852 }
853 _ => None,
854 }
855}
856
857pub fn resolve_module_code_unit(
865 python: &dyn PythonSource,
866 importer: Option<&ProjectFile>,
867 module_fq: &str,
868) -> Option<CodeUnit> {
869 let indexed = match python.path_module_fqn(module_fq) {
870 Some(units) => indexed_module_for_importer(python.project(), importer, units),
873 None => indexed_module_for_importer(
874 python.project(),
875 importer,
876 python.definition_fqn(module_fq),
877 ),
878 };
879 match indexed {
880 IndexedModuleLookup::Unique(module) => Some(module),
881 IndexedModuleLookup::Ambiguous => None,
885 IndexedModuleLookup::Missing => {
886 let importer = importer?;
887 module_under_source_root(&python.module_spellings(), importer, module_fq)
888 }
889 }
890}
891
892enum IndexedModuleLookup {
901 Unique(CodeUnit),
902 Ambiguous,
903 Missing,
904}
905
906fn indexed_module_for_importer(
907 project: &dyn Project,
908 importer: Option<&ProjectFile>,
909 units: Vec<CodeUnit>,
910) -> IndexedModuleLookup {
911 let mut modules = units
912 .into_iter()
913 .filter(CodeUnit::is_module)
914 .collect::<Vec<_>>();
915 if let Some(importer) = importer {
916 retain_modules_for_importer(project, importer, &mut modules, CodeUnit::source);
917 }
918 match modules.len() {
919 0 => IndexedModuleLookup::Missing,
920 1 => IndexedModuleLookup::Unique(modules.remove(0)),
921 _ => IndexedModuleLookup::Ambiguous,
922 }
923}
924
925pub fn python_module_file_under_source_root(
954 spellings: &PythonModuleSpellings,
955 importer: &ProjectFile,
956 module_fq: &str,
957) -> Option<ProjectFile> {
958 let mut candidates = spellings.suffix_candidates(module_fq);
959 retain_nearest_owning_root(importer, &mut candidates, |candidate| {
960 candidate.source_root.clone()
961 });
962 candidates.retain(|candidate| spellings.establishes_source_root(candidate));
963 let [candidate] = candidates.as_slice() else {
964 return None;
965 };
966 debug_assert_eq!(
967 importer.root(),
968 candidate.file.root(),
969 "one workspace's Python files share one project root"
970 );
971 Some(candidate.file.clone())
972}
973
974fn module_under_source_root(
979 spellings: &PythonModuleSpellings,
980 importer: &ProjectFile,
981 module_fq: &str,
982) -> Option<CodeUnit> {
983 let file = python_module_file_under_source_root(spellings, importer, module_fq)?;
984 module_code_unit(spellings.project()?, &file)
985}
986
987pub fn resolve_module_code_units_batch(
993 python: &dyn PythonSource,
994 importer: Option<&ProjectFile>,
995 module_fqs: &[String],
996) -> Vec<Option<CodeUnit>> {
997 let path_results = python.path_module_fqns_batch(module_fqs);
998 let mut lookups: Vec<IndexedModuleLookup> = Vec::with_capacity(module_fqs.len());
999 let mut needs_definition_fallback = Vec::new();
1000 for (i, units) in path_results.into_iter().enumerate() {
1001 match units {
1002 Some(units) => lookups.push(indexed_module_for_importer(
1003 python.project(),
1004 importer,
1005 units,
1006 )),
1007 None => {
1008 needs_definition_fallback.push(i);
1009 lookups.push(IndexedModuleLookup::Missing);
1011 }
1012 }
1013 }
1014 for i in needs_definition_fallback {
1015 lookups[i] = indexed_module_for_importer(
1016 python.project(),
1017 importer,
1018 python.definition_fqn(&module_fqs[i]),
1019 );
1020 }
1021 let mut results: Vec<Option<CodeUnit>> = vec![None; module_fqs.len()];
1022 let mut needs_source_root = Vec::new();
1023 for (i, lookup) in lookups.into_iter().enumerate() {
1024 match lookup {
1025 IndexedModuleLookup::Unique(module) => results[i] = Some(module),
1026 IndexedModuleLookup::Ambiguous => {}
1027 IndexedModuleLookup::Missing => needs_source_root.push(i),
1028 }
1029 }
1030 if let Some(importer) = importer
1031 && !needs_source_root.is_empty()
1032 {
1033 let spellings = python.module_spellings();
1036 for i in needs_source_root {
1037 results[i] = module_under_source_root(&spellings, importer, &module_fqs[i]);
1038 }
1039 }
1040 results
1041}
1042
1043pub fn compute_export_index_of(
1044 python: &dyn PythonSource,
1045 token: QueryToken<'_>,
1046 file: &ProjectFile,
1047) -> ExportIndex {
1048 let mut index = ExportIndex::empty();
1049 let mut events = Vec::new();
1050 let declarations = python.top_level_declarations(file);
1051 collect_local_export_events(
1052 declarations.iter(),
1053 |code_unit| {
1054 python
1055 .ranges(code_unit)
1056 .iter()
1057 .map(|range| range.start_byte)
1058 .max()
1059 .unwrap_or(usize::MAX)
1060 },
1061 &mut events,
1062 );
1063
1064 let imports = python.import_info_of(token, file);
1065 collect_reexport_events_from_imports(python, file, &imports, &mut events, &mut index);
1066
1067 finish_export_index(events, index)
1068}
1069
1070pub fn export_index_from_file_facts(
1071 python: &dyn PythonSource,
1072 file: &ProjectFile,
1073 facts: &dyn IndexedFileFacts,
1074 module_name: &str,
1075) -> ExportIndex {
1076 let mut index = ExportIndex::empty();
1077 let mut events = Vec::new();
1078 collect_local_export_events(
1079 facts.top_level_declarations().iter(),
1080 |code_unit| {
1081 facts
1082 .declaration_ranges(code_unit)
1083 .into_iter()
1084 .flatten()
1085 .map(|range| range.start_byte)
1086 .max()
1087 .unwrap_or(usize::MAX)
1088 },
1089 &mut events,
1090 );
1091
1092 if !facts
1093 .top_level_declarations()
1094 .iter()
1095 .any(CodeUnit::is_module)
1096 && let Some(identifier) = module_name.rsplit('.').next()
1097 && !identifier.is_empty()
1098 && !identifier.starts_with('_')
1099 {
1100 events.push((
1101 0,
1102 identifier.to_string(),
1103 ExportEntry::Local {
1104 local_name: identifier.to_string(),
1105 },
1106 ));
1107 }
1108
1109 collect_reexport_events_from_imports(python, file, facts.imports(), &mut events, &mut index);
1110
1111 finish_export_index(events, index)
1112}
1113
1114fn collect_local_export_events<'a>(
1115 declarations: impl IntoIterator<Item = &'a CodeUnit>,
1116 mut start_byte: impl FnMut(&CodeUnit) -> usize,
1117 events: &mut Vec<(usize, String, ExportEntry)>,
1118) -> HashSet<String> {
1119 let mut local_names = HashSet::default();
1120 for code_unit in declarations {
1121 let identifier = code_unit.identifier().trim();
1122 if identifier.is_empty() {
1123 continue;
1124 }
1125 local_names.insert(identifier.to_string());
1126 events.push((
1127 start_byte(code_unit),
1128 identifier.to_string(),
1129 ExportEntry::Local {
1130 local_name: identifier.to_string(),
1131 },
1132 ));
1133 }
1134 local_names
1135}
1136
1137fn finish_export_index(
1138 mut events: Vec<(usize, String, ExportEntry)>,
1139 mut index: ExportIndex,
1140) -> ExportIndex {
1141 events.sort_by_key(|(start_byte, _, _)| *start_byte);
1142 for (_, exported_name, entry) in events {
1143 index.exports_by_name.insert(exported_name, entry);
1144 }
1145 index
1146}
1147
1148fn collect_reexport_events_from_imports(
1149 python: &dyn PythonSource,
1150 file: &ProjectFile,
1151 imports: &[ImportInfo],
1152 events: &mut Vec<(usize, String, ExportEntry)>,
1153 index: &mut ExportIndex,
1154) {
1155 for import in imports {
1156 if import
1157 .path
1158 .as_ref()
1159 .is_none_or(|path| !path.lexical_scopes.is_empty())
1160 {
1161 continue;
1162 }
1163 record_single_reexport_event(python, file, import, events, index);
1164 }
1165}
1166
1167fn record_single_reexport_event(
1168 python: &dyn PythonSource,
1169 file: &ProjectFile,
1170 import: &ImportInfo,
1171 events: &mut Vec<(usize, String, ExportEntry)>,
1172 index: &mut ExportIndex,
1173) {
1174 let Some(PythonImportDetails::FromImport {
1175 module,
1176 name,
1177 alias,
1178 wildcard,
1179 }) = python_import_details(import)
1180 else {
1181 return;
1182 };
1183 let start_byte = import
1184 .path
1185 .as_ref()
1186 .map(|path| path.declaration_start_byte)
1187 .unwrap_or(usize::MAX);
1188 let resolved_module = if module.starts_with('.') {
1189 resolve_python_relative_module(python.project(), file, &module)
1190 } else {
1191 Some(module.clone())
1192 };
1193 let Some(resolved_module) = resolved_module else {
1194 return;
1195 };
1196
1197 if wildcard {
1198 index.reexport_stars.push(ReexportStar {
1199 module_specifier: resolved_module,
1200 });
1201 return;
1202 }
1203 let exported_name = alias.unwrap_or(name.clone());
1204 let module_candidate = format!("{resolved_module}.{name}");
1210 if resolve_module_code_unit(python, Some(file), &module_candidate).is_some() {
1211 events.push((
1212 start_byte,
1213 exported_name,
1214 ExportEntry::ReexportedModule {
1215 module_specifier: module_candidate,
1216 },
1217 ));
1218 return;
1219 }
1220 events.push((
1221 start_byte,
1222 exported_name,
1223 ExportEntry::ReexportedNamed {
1224 module_specifier: resolved_module,
1225 imported_name: name,
1226 },
1227 ));
1228}
1229
1230pub fn import_binder_from_imports(
1231 python: &dyn PythonSource,
1232 file: &ProjectFile,
1233 imports: &[ImportInfo],
1234) -> ImportBinder {
1235 let mut binder = ImportBinder::empty();
1236
1237 for (local_name, binding) in import_bindings_from_imports(python, file, imports) {
1238 binder.bindings.insert(local_name, binding);
1239 }
1240
1241 binder
1242}
1243
1244pub fn import_bindings_from_imports(
1249 python: &dyn PythonSource,
1250 file: &ProjectFile,
1251 imports: &[ImportInfo],
1252) -> Vec<(String, ImportBinding)> {
1253 let mut bindings = Vec::new();
1254
1255 for import in imports {
1256 let Some(details) = python_import_details(import) else {
1257 continue;
1258 };
1259 match details {
1260 PythonImportDetails::Import { module, alias } => {
1261 let local_name = python_namespace_binding_name(import, alias.as_deref(), &module);
1262 let module_specifier =
1263 python_namespace_binding_module(import, alias.as_deref(), &module);
1264 bindings.push((
1265 local_name,
1266 ImportBinding {
1267 module_specifier,
1268 namespace_imported_module: Some(module),
1269 kind: ImportKind::Namespace,
1270 imported_name: None,
1271 },
1272 ));
1273 }
1274 PythonImportDetails::FromImport {
1275 module,
1276 name,
1277 wildcard,
1278 ..
1279 } => {
1280 let resolved_module = if module.starts_with('.') {
1281 resolve_python_relative_module(python.project(), file, &module)
1282 } else {
1283 Some(module.clone())
1284 };
1285 let Some(resolved_module) = resolved_module else {
1286 continue;
1287 };
1288 if wildcard {
1289 bindings.extend(
1294 public_declarations_in_module(python, file, &resolved_module)
1295 .into_iter()
1296 .map(|declaration| {
1297 let name = declaration.identifier().to_string();
1298 (
1299 name.clone(),
1300 ImportBinding {
1301 module_specifier: resolved_module.clone(),
1302 namespace_imported_module: None,
1303 kind: ImportKind::Named,
1304 imported_name: Some(name),
1305 },
1306 )
1307 }),
1308 );
1309 continue;
1310 }
1311 let local_name = import
1316 .local_name()
1317 .map(str::to_string)
1318 .unwrap_or_else(|| name.clone());
1319 let module_candidate = format!("{resolved_module}.{name}");
1320 let exported = resolve_exported_fqn(python, Some(file), &module_candidate);
1321 let exported_module =
1322 if !exported.is_empty() && exported.iter().all(CodeUnit::is_module) {
1323 let mut paths = exported.iter().map(CodeUnit::fq_segment_texts);
1324 let path = paths
1325 .next()
1326 .expect("non-empty exported modules have a first path");
1327 paths
1328 .all(|candidate| candidate == path)
1329 .then(|| path.join("."))
1330 } else {
1331 None
1332 };
1333 let implicit_submodule = if exported.is_empty()
1334 && resolve_module_code_unit(python, Some(file), &module_candidate).is_some()
1335 {
1336 Some(module_candidate)
1337 } else {
1338 None
1339 };
1340 if let Some(module_specifier) = exported_module.or(implicit_submodule) {
1341 bindings.push((
1342 local_name,
1343 ImportBinding {
1344 module_specifier,
1345 namespace_imported_module: None,
1346 kind: ImportKind::Namespace,
1347 imported_name: None,
1348 },
1349 ));
1350 continue;
1351 }
1352 bindings.push((
1353 local_name,
1354 ImportBinding {
1355 module_specifier: resolved_module,
1356 namespace_imported_module: None,
1357 kind: ImportKind::Named,
1358 imported_name: Some(name),
1359 },
1360 ));
1361 }
1362 }
1363 }
1364
1365 bindings
1366}
1367
1368pub fn public_declarations_in_module(
1369 python: &dyn PythonSource,
1370 importer: &ProjectFile,
1371 module_fq: &str,
1372) -> Vec<CodeUnit> {
1373 let Some(module_code_unit) = resolve_module_code_unit(python, Some(importer), module_fq) else {
1374 return Vec::new();
1375 };
1376 python
1377 .direct_children(&module_code_unit)
1378 .into_iter()
1379 .filter(|code_unit| !code_unit.identifier().starts_with('_'))
1380 .collect()
1381}
1382
1383pub fn resolve_base_class(
1384 python: &dyn PythonSource,
1385 token: QueryToken<'_>,
1386 code_unit: &CodeUnit,
1387 raw: &str,
1388) -> Option<CodeUnit> {
1389 let trimmed = raw.trim();
1390 if trimmed.is_empty() {
1391 return None;
1392 }
1393
1394 let binder = python.import_binder_of(code_unit.source());
1395 if let Some((head, tail)) = trimmed.split_once('.') {
1396 if let Some(binding) = binder.bindings.get(head)
1397 && binding.kind == ImportKind::Namespace
1398 {
1399 let fq_name = format!("{}.{}", binding.module_specifier, tail);
1400 return python.definitions(&fq_name).next();
1401 }
1402 return python.definitions(trimmed).next();
1403 }
1404
1405 if let Some(binding) = binder.bindings.get(trimmed) {
1406 match binding.kind {
1407 ImportKind::Namespace => {
1408 return resolve_module_code_unit(
1409 python,
1410 Some(code_unit.source()),
1411 &binding.module_specifier,
1412 );
1413 }
1414 ImportKind::Named => {
1415 let imported_name = binding.imported_name.as_ref()?;
1416 let fqn = format!("{}.{}", binding.module_specifier, imported_name);
1417 return resolve_exported_fqn(python, Some(code_unit.source()), &fqn)
1418 .into_iter()
1419 .next()
1420 .or_else(|| python.definitions(&fqn).next());
1421 }
1422 _ => {}
1423 }
1424 }
1425
1426 if python
1427 .import_info_of(token, code_unit.source())
1428 .iter()
1429 .any(|import| import.is_wildcard)
1430 && let Some(imported) =
1431 resolve_import_bindings(python, token, code_unit.source()).get(trimmed)
1432 {
1433 return Some(imported.clone());
1434 }
1435
1436 if let Some(local) = resolve_preceding_local_class_base(python, token, code_unit, trimmed) {
1437 return Some(local);
1438 }
1439
1440 let local_fq_name = format!("{}.{}", code_unit.package_name(), trimmed);
1441 python
1442 .definitions(&local_fq_name)
1443 .next()
1444 .or_else(|| python.definitions(trimmed).next())
1445}
1446
1447fn class_definition_in_statement(node: Node<'_>) -> Option<Node<'_>> {
1448 if node.kind() == "class_definition" {
1449 return Some(node);
1450 }
1451 (node.kind() == "decorated_definition")
1452 .then(|| {
1453 let mut cursor = node.walk();
1454 node.named_children(&mut cursor)
1455 .find(|child| child.kind() == "class_definition")
1456 })
1457 .flatten()
1458}
1459
1460fn declaration_statement(node: Node<'_>) -> Option<Node<'_>> {
1461 match node.kind() {
1462 "decorated_definition" if class_definition_in_statement(node).is_some() => Some(node),
1463 "class_definition" => Some(
1464 node.parent()
1465 .filter(|parent| parent.kind() == "decorated_definition")
1466 .unwrap_or(node),
1467 ),
1468 _ => None,
1469 }
1470}
1471
1472fn scope_statement_binds_name(node: Node<'_>, source: &str, name: &str) -> bool {
1473 let mut stack = vec![node];
1474 while let Some(node) = stack.pop() {
1475 if node.kind() == "wildcard_import" {
1476 return true;
1477 }
1478 if matches!(node.kind(), "global_statement" | "nonlocal_statement") {
1479 let mut cursor = node.walk();
1480 if node.named_children(&mut cursor).any(|child| {
1481 child.kind() == "identifier"
1482 && child.utf8_text(source.as_bytes()).ok() == Some(name)
1483 }) {
1484 return true;
1485 }
1486 }
1487 let Some(bindings) = python_direct_scope_bindings_bounded(node, source, || true) else {
1488 return true;
1489 };
1490 if bindings
1491 .iter()
1492 .any(|binding| binding.declaration.utf8_text(source.as_bytes()).ok() == Some(name))
1493 {
1494 return true;
1495 }
1496 stack.extend(python_executable_scope_children(node));
1497 }
1498 false
1499}
1500
1501fn resolve_preceding_local_class_base(
1502 python: &dyn PythonSource,
1503 token: QueryToken<'_>,
1504 code_unit: &CodeUnit,
1505 raw: &str,
1506) -> Option<CodeUnit> {
1507 let parent = python.parent_of(code_unit)?;
1508 if !parent.is_function() && !parent.is_class() {
1509 return None;
1510 }
1511 let prepared = python.prepared_syntax(token, code_unit.source())?;
1512 let source = prepared.source();
1513 let current_node = prepared.declaration_node(code_unit)?;
1514 let current_statement = declaration_statement(current_node)?;
1515 let current_class = class_definition_in_statement(current_statement)?;
1516 let base_list = current_class.child_by_field_name("superclasses")?;
1517 let mut base_cursor = base_list.walk();
1518 if !base_list.named_children(&mut base_cursor).any(|base| {
1519 let origin = python_base_origin_node(base);
1520 origin.kind() == "identifier" && origin.utf8_text(source.as_bytes()).ok() == Some(raw)
1521 }) {
1522 return None;
1523 }
1524 let block = current_statement.parent()?;
1525 let mut block_cursor = block.walk();
1526 let siblings = block.named_children(&mut block_cursor).collect::<Vec<_>>();
1527 let current_index = siblings
1528 .iter()
1529 .position(|sibling| sibling.id() == current_statement.id())?;
1530
1531 let mut preceding: Option<(usize, CodeUnit, Node<'_>)> = None;
1532 for candidate in prepared.direct_children(&parent) {
1533 if !candidate.is_class()
1534 || candidate.source() != code_unit.source()
1535 || candidate.identifier() != raw
1536 {
1537 continue;
1538 }
1539 let Some(candidate_node) = prepared.declaration_node(candidate) else {
1540 continue;
1541 };
1542 let Some(candidate_statement) = declaration_statement(candidate_node) else {
1543 continue;
1544 };
1545 if candidate_statement.parent().map(|node| node.id()) != Some(block.id())
1546 || candidate_statement.start_byte() >= current_statement.start_byte()
1547 {
1548 continue;
1549 }
1550 if preceding
1551 .as_ref()
1552 .is_none_or(|(start, _, _)| candidate_statement.start_byte() > *start)
1553 {
1554 preceding = Some((
1555 candidate_statement.start_byte(),
1556 candidate.clone(),
1557 candidate_statement,
1558 ));
1559 }
1560 }
1561 let (_, candidate, candidate_statement) = preceding?;
1562 let candidate_index = siblings
1563 .iter()
1564 .position(|sibling| sibling.id() == candidate_statement.id())?;
1565 if candidate_index >= current_index
1566 || siblings[candidate_index + 1..current_index]
1567 .iter()
1568 .any(|statement| scope_statement_binds_name(*statement, source, raw))
1569 {
1570 return None;
1571 }
1572 Some(candidate)
1573}
1574
1575pub fn render_skeleton_recursive(
1576 index: &dyn CodeUnitIndex,
1577 code_unit: &CodeUnit,
1578 indent: &str,
1579 header_only: bool,
1580 out: &mut String,
1581) {
1582 enum Work {
1583 Declaration(CodeUnit, String),
1584 Elision(String),
1585 }
1586 let mut stack = vec![Work::Declaration(code_unit.clone(), indent.to_string())];
1587 while let Some(work) = stack.pop() {
1588 let Work::Declaration(unit, indent) = work else {
1589 let Work::Elision(indent) = work else {
1590 unreachable!()
1591 };
1592 out.push_str(&indent);
1593 out.push_str("[...]\n");
1594 continue;
1595 };
1596 if !unit.is_module()
1597 && !unit.is_file_scope()
1598 && let Some(signature) = index.signatures(&unit).first()
1599 {
1600 let signature = if unit.is_field() {
1603 signature.strip_suffix(" ...").unwrap_or(signature)
1604 } else {
1605 signature
1606 };
1607 for line in signature.lines() {
1608 out.push_str(&indent);
1609 out.push_str(line);
1610 out.push('\n');
1611 }
1612 }
1613 let children = index.direct_children(&unit);
1614 let child_indent = format!("{indent} ");
1615 if header_only && children.iter().any(|child| !child.is_field()) {
1616 stack.push(Work::Elision(child_indent.clone()));
1617 }
1618 for child in children
1619 .into_iter()
1620 .rev()
1621 .filter(|child| !header_only || child.is_field())
1622 {
1623 stack.push(Work::Declaration(child, child_indent.clone()));
1624 }
1625 }
1626}