1use crate::bindings::python_direct_scope_bindings_bounded;
2use crate::imports::python_import_infos_from_node;
3use crate::syntax::{
4 PythonOverloadDecoratorBindings, expression_name_node, python_plain_string_literal,
5};
6use brokk_bifrost_core::analyzer::fq_name::{FqName, SegmentId, SegmentKind, segment_interner};
7use brokk_bifrost_core::analyzer::model::{
8 CodeUnitType, DispatchExtensibility, ParameterMetadata, SignatureMetadata,
9 StructuredImportPathKind,
10};
11use brokk_bifrost_core::analyzer::parsed_file::ParsedFile;
12use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
13use brokk_bifrost_core::analyzer::{CodeUnit, ProjectFile};
14use brokk_bifrost_core::hash::{HashMap, HashSet};
15use brokk_bifrost_core::path_normalization::NormalizePath;
16use brokk_bifrost_core::text_utils::{compute_line_starts, find_line_index_for_offset};
17use std::path::{Path, PathBuf};
18use tree_sitter::{Node, Parser, Tree};
19
20fn py_segment(text: &str, kind: SegmentKind) -> SegmentId {
22 segment_interner().intern(text, kind)
23}
24
25pub fn python_module_fq(file: &ProjectFile) -> FqName {
38 python_module_fq_from_components(&python_module_components(file))
39}
40
41fn python_module_fq_from_components(components: &[String]) -> FqName {
42 let mut fq = FqName::new();
43 for component in components {
44 fq.push(py_segment(component, SegmentKind::Package));
45 }
46 fq
47}
48
49fn python_module_components(file: &ProjectFile) -> Vec<String> {
50 let mut components = python_package_components_for_file(file);
51 let module_name = file
52 .rel_path()
53 .file_stem()
54 .and_then(|stem| stem.to_str())
55 .unwrap_or_default();
56 if module_name != "__init__" || components.is_empty() {
57 components.push(module_name.to_string());
58 }
59 components
60}
61
62fn python_package_components_for_file(file: &ProjectFile) -> Vec<String> {
63 let Some(parent_rel) = file.rel_path().parent() else {
64 return Vec::new();
65 };
66 if parent_rel.as_os_str().is_empty() {
67 return Vec::new();
68 }
69
70 if let Some(import_root_rel) = python_configured_import_root(file, parent_rel)
71 && let Ok(relative_package) = parent_rel.strip_prefix(import_root_rel)
72 {
73 return path_components(relative_package);
74 }
75
76 let mut effective_package_root_rel: Option<&Path> = None;
77 let mut current_rel = Some(parent_rel);
78 while let Some(path) = current_rel {
79 if file.root().join(path).join("__init__.py").exists() {
80 effective_package_root_rel = Some(path);
81 }
82 current_rel = path.parent();
83 }
84
85 let relative_package = match effective_package_root_rel {
86 Some(package_root_rel) => package_root_rel
87 .parent()
88 .and_then(|import_root_rel| parent_rel.strip_prefix(import_root_rel).ok())
89 .unwrap_or(parent_rel),
90 None => parent_rel,
91 };
92 path_components(relative_package)
93}
94
95fn python_configured_import_root(file: &ProjectFile, parent_rel: &Path) -> Option<PathBuf> {
101 let mut manifest_dir_rel = Some(parent_rel);
102 while let Some(directory) = manifest_dir_rel {
103 let manifest_dir = file.root().join(directory);
104 let mut roots = setuptools_where_entries(&manifest_dir.join("pyproject.toml"))
105 .iter()
106 .map(|entry| manifest_dir.join(entry).normalize())
107 .filter_map(|root| root.strip_prefix(file.root()).ok().map(Path::to_path_buf))
108 .filter(|root| parent_rel.starts_with(root))
109 .collect::<Vec<_>>();
110 roots.sort_by_key(|root| root.components().count());
111 if let Some(root) = roots.pop() {
112 return Some(root);
113 }
114 if let Some(package_dir) = setuptools_setup_py_import_root(&manifest_dir.join("setup.py")) {
115 let root = manifest_dir.join(package_dir).normalize();
116 if let Ok(root) = root.strip_prefix(file.root())
117 && parent_rel.starts_with(root)
118 {
119 return Some(root.to_path_buf());
120 }
121 }
122 manifest_dir_rel = directory.parent();
123 }
124 None
125}
126
127#[derive(Clone, Copy, PartialEq, Eq)]
128struct FileStamp {
129 len: u64,
130 modified: Option<std::time::SystemTime>,
131}
132
133fn file_stamp(path: &Path) -> Option<FileStamp> {
134 let metadata = std::fs::metadata(path).ok()?;
135 Some(FileStamp {
136 len: metadata.len(),
137 modified: metadata.modified().ok(),
138 })
139}
140
141struct ManifestWhereEntries {
146 stamp: FileStamp,
147 entries: Vec<String>,
148}
149
150fn setuptools_where_entries(manifest: &Path) -> Vec<String> {
161 static MEMO: std::sync::OnceLock<
162 std::sync::RwLock<std::collections::HashMap<PathBuf, ManifestWhereEntries>>,
163 > = std::sync::OnceLock::new();
164 let memo = MEMO.get_or_init(Default::default);
165
166 let Some(stamp) = file_stamp(manifest) else {
167 return Vec::new();
168 };
169 if let Some(cached) = memo.read().expect("manifest memo").get(manifest)
170 && cached.stamp == stamp
171 {
172 return cached.entries.clone();
173 }
174
175 let entries = parse_setuptools_where_entries(manifest);
176 memo.write().expect("manifest memo").insert(
177 manifest.to_path_buf(),
178 ManifestWhereEntries {
179 stamp,
180 entries: entries.clone(),
181 },
182 );
183 entries
184}
185
186struct SetupPyImportRoot {
190 stamp: FileStamp,
191 package_dir: Option<PathBuf>,
192}
193
194fn setuptools_setup_py_import_root(setup_py: &Path) -> Option<PathBuf> {
204 static MEMO: std::sync::OnceLock<
205 std::sync::RwLock<std::collections::HashMap<PathBuf, SetupPyImportRoot>>,
206 > = std::sync::OnceLock::new();
207 let memo = MEMO.get_or_init(Default::default);
208
209 let stamp = file_stamp(setup_py)?;
210 if let Some(cached) = memo.read().expect("setup.py memo").get(setup_py)
211 && cached.stamp == stamp
212 {
213 return cached.package_dir.clone();
214 }
215
216 let package_dir = parse_setuptools_setup_py_import_root(setup_py);
217 memo.write().expect("setup.py memo").insert(
218 setup_py.to_path_buf(),
219 SetupPyImportRoot {
220 stamp,
221 package_dir: package_dir.clone(),
222 },
223 );
224 package_dir
225}
226
227fn parse_setuptools_setup_py_import_root(setup_py: &Path) -> Option<PathBuf> {
228 let source = std::fs::read_to_string(setup_py).ok()?;
229 let tree = parse_python_tree(&source)?;
230 let root = tree.root_node();
231 if root.has_error() {
232 return None;
233 }
234 let mut setup_bindings: HashMap<Vec<String>, String> = HashMap::default();
235 let mut import_root = None;
236 let mut cursor = root.walk();
237
238 for statement in root.named_children(&mut cursor) {
239 if matches!(
240 statement.kind(),
241 "import_statement" | "import_from_statement"
242 ) {
243 for binding in setup_py_bound_names(statement, &source) {
244 setup_bindings.retain(|path, _| path.first() != Some(&binding));
245 }
246 for import in python_import_infos_from_node(statement, &source) {
247 if import.is_wildcard {
248 setup_bindings.clear();
249 continue;
250 }
251 let Some(path) = import.path else { continue };
252 let segments = path.segments.iter().map(String::as_str).collect::<Vec<_>>();
253 match path.kind {
254 Some(StructuredImportPathKind::ImportFrom) => {
255 let function_name = match segments.as_slice() {
256 ["setuptools", "setup"] | ["distutils", "core", "setup"] => "setup",
257 ["setuptools", "find_packages"] => "find_packages",
258 ["setuptools", "find_namespace_packages"] => "find_namespace_packages",
259 _ => continue,
260 };
261 setup_bindings.insert(
262 vec![import.identifier.expect("imported function binds a name")],
263 function_name.to_string(),
264 );
265 }
266 Some(StructuredImportPathKind::Namespace) => {
267 let function_names = match segments.as_slice() {
268 ["setuptools"] => {
269 ["setup", "find_packages", "find_namespace_packages"].as_slice()
270 }
271 ["distutils", "core"] => ["setup"].as_slice(),
272 _ => continue,
273 };
274 let binding_prefix = import
275 .alias
276 .map(|alias| vec![alias])
277 .unwrap_or_else(|| path.segments.clone());
278 for function_name in function_names {
279 let mut callable = binding_prefix.clone();
280 callable.push((*function_name).to_string());
281 setup_bindings.insert(callable, (*function_name).to_string());
282 }
283 }
284 _ => continue,
285 }
286 }
287 continue;
288 }
289
290 if statement.kind() == "expression_statement"
291 && statement.named_child_count() == 1
292 && let Some(call) = statement.named_child(0)
293 && call.kind() == "call"
294 && setup_py_call_imported_function(call, &source, &setup_bindings)
295 .is_some_and(|function| function == "setup")
296 {
297 let candidate = setup_py_import_root_from_call(call, &source, &setup_bindings)?;
298 if import_root
299 .replace(candidate.clone())
300 .is_some_and(|root| root != candidate)
301 {
302 return None;
303 }
304 }
305
306 for binding in setup_py_bound_names(statement, &source) {
307 setup_bindings.retain(|path, _| path.first() != Some(&binding));
308 }
309 }
310 import_root
311}
312
313fn setup_py_bound_names(statement: Node<'_>, source: &str) -> Vec<String> {
320 let mut names = Vec::new();
321 let mut pending = vec![statement];
322 while let Some(node) = pending.pop() {
323 for binding in python_direct_scope_bindings_bounded(node, source, || true)
324 .expect("unbounded setup.py binding walk")
325 {
326 let name = py_node_text(binding.declaration, source).trim();
327 if !name.is_empty() {
328 names.push(name.to_string());
329 }
330 }
331 let excluded_body = matches!(
332 node.kind(),
333 "function_definition" | "class_definition" | "lambda"
334 )
335 .then(|| node.child_by_field_name("body").map(|body| body.id()))
336 .flatten();
337 let mut cursor = node.walk();
338 pending.extend(
339 node.named_children(&mut cursor)
340 .filter(|child| Some(child.id()) != excluded_body),
341 );
342 }
343 names
344}
345
346fn setup_py_call_imported_function<'a>(
347 call: Node<'_>,
348 source: &str,
349 setup_bindings: &'a HashMap<Vec<String>, String>,
350) -> Option<&'a str> {
351 let mut function = call.child_by_field_name("function")?;
352 let mut path = Vec::new();
353 while function.kind() == "attribute" {
354 let attribute = function.child_by_field_name("attribute")?;
355 path.push(py_node_text(attribute, source).to_string());
356 let object = function.child_by_field_name("object")?;
357 function = object;
358 }
359 if function.kind() != "identifier" {
360 return None;
361 }
362 path.push(py_node_text(function, source).to_string());
363 path.reverse();
364 setup_bindings.get(&path).map(String::as_str)
365}
366
367fn setup_py_import_root_from_call(
368 call: Node<'_>,
369 source: &str,
370 setup_bindings: &HashMap<Vec<String>, String>,
371) -> Option<PathBuf> {
372 let arguments = call.child_by_field_name("arguments")?;
373 if arguments.kind() != "argument_list" {
374 return None;
375 }
376
377 let mut packages = None;
378 let mut package_dir = None;
379 let mut cursor = arguments.walk();
380 for argument in arguments.named_children(&mut cursor) {
381 if argument.kind() == "comment" {
382 continue;
383 }
384 if argument.kind() != "keyword_argument" {
386 return None;
387 }
388 let name = argument.child_by_field_name("name")?;
389 let value = argument.child_by_field_name("value")?;
390 match py_node_text(name, source).trim() {
391 "packages" if packages.is_none() => packages = Some(value),
392 "packages" => return None,
393 "package_dir" if package_dir.is_none() => package_dir = Some(value),
394 "package_dir" => return None,
395 _ => {}
396 }
397 }
398
399 let packages = packages?;
400 if let Some(package_dir) = package_dir {
401 return setup_py_static_package_dir(package_dir, source);
402 }
403 if setup_py_nonempty_literal_packages(packages, source) {
404 return Some(PathBuf::new());
405 }
406 setup_py_discovery_root_from_call(packages, source, setup_bindings)
407}
408
409fn setup_py_discovery_root_from_call(
410 call: Node<'_>,
411 source: &str,
412 setup_bindings: &HashMap<Vec<String>, String>,
413) -> Option<PathBuf> {
414 let function = setup_py_call_imported_function(call, source, setup_bindings)?;
415 if !matches!(function, "find_packages" | "find_namespace_packages") {
416 return None;
417 }
418 let arguments = call.child_by_field_name("arguments")?;
419 if arguments.kind() != "argument_list" {
420 return None;
421 }
422
423 let mut where_value = None;
424 let mut seen_exclude = false;
425 let mut seen_include = false;
426 let mut positional_index = 0;
427 let mut cursor = arguments.walk();
428 for argument in arguments.named_children(&mut cursor) {
429 if argument.kind() == "comment" {
430 continue;
431 }
432 if matches!(argument.kind(), "list_splat" | "dictionary_splat") {
433 return None;
434 }
435 if argument.kind() == "keyword_argument" {
436 let name = py_node_text(argument.child_by_field_name("name")?, source).trim();
437 let value = argument.child_by_field_name("value")?;
438 match name {
439 "where" if where_value.is_none() => where_value = Some(value),
440 "where" => return None,
441 "exclude" if !seen_exclude => seen_exclude = true,
442 "exclude" => return None,
443 "include" if !seen_include => seen_include = true,
444 "include" => return None,
445 _ => return None,
446 }
447 continue;
448 }
449
450 let slot = positional_index;
451 positional_index += 1;
452 match slot {
453 0 if where_value.is_none() => where_value = Some(argument),
454 0 => return None,
455 1 if !seen_exclude => seen_exclude = true,
456 1 => return None,
457 2 if !seen_include => seen_include = true,
458 2 => return None,
459 _ => return None,
460 }
461 }
462
463 where_value
464 .map(|value| python_plain_string_literal(value, source))
465 .unwrap_or(Some(""))
466 .map(PathBuf::from)
467}
468
469fn setup_py_nonempty_literal_packages(value: Node<'_>, source: &str) -> bool {
470 let value = setup_py_unwrap_parenthesized(value);
471 if !matches!(value.kind(), "list" | "set" | "tuple") {
472 return false;
473 }
474 let mut cursor = value.walk();
475 let mut nonempty = false;
476 for element in value.named_children(&mut cursor) {
477 if element.kind() == "comment" {
478 continue;
479 }
480 let Some(package) = python_plain_string_literal(element, source) else {
481 return false;
482 };
483 if package.is_empty() {
484 return false;
485 }
486 nonempty = true;
487 }
488 nonempty
489}
490
491fn setup_py_static_package_dir(value: Node<'_>, source: &str) -> Option<PathBuf> {
492 let value = setup_py_unwrap_parenthesized(value);
493 if value.kind() != "dictionary" {
494 return None;
495 }
496 let mut cursor = value.walk();
497 let mut pairs = value
498 .named_children(&mut cursor)
499 .filter(|node| node.kind() != "comment");
500 let Some(pair) = pairs.next() else {
501 return Some(PathBuf::new());
502 };
503 if pairs.next().is_some() || pair.kind() != "pair" {
504 return None;
505 }
506 let key = python_plain_string_literal(pair.child_by_field_name("key")?, source)?;
507 if !key.is_empty() {
508 return None;
509 }
510 let root = python_plain_string_literal(pair.child_by_field_name("value")?, source)?;
511 let root = PathBuf::from(root);
512 (!root.is_absolute()).then_some(root)
513}
514
515fn setup_py_unwrap_parenthesized(mut node: Node<'_>) -> Node<'_> {
516 while node.kind() == "parenthesized_expression" && node.named_child_count() == 1 {
517 node = node.named_child(0).expect("parenthesized expression child");
518 }
519 node
520}
521
522fn parse_setuptools_where_entries(manifest: &Path) -> Vec<String> {
523 let Ok(source) = std::fs::read_to_string(manifest) else {
524 return Vec::new();
525 };
526 let Ok(document) = source.parse::<toml::Value>() else {
527 return Vec::new();
528 };
529 document
530 .get("tool")
531 .and_then(|tool| tool.get("setuptools"))
532 .and_then(|setuptools| setuptools.get("packages"))
533 .and_then(|packages| packages.get("find"))
534 .and_then(|find| find.get("where"))
535 .and_then(toml::Value::as_array)
536 .map(|entries| {
537 entries
538 .iter()
539 .filter_map(toml::Value::as_str)
540 .map(str::to_string)
541 .collect()
542 })
543 .unwrap_or_default()
544}
545
546fn path_components(path: &Path) -> Vec<String> {
547 path.components()
548 .map(|component| component.as_os_str().to_string_lossy().to_string())
549 .filter(|component| !component.is_empty())
550 .collect()
551}
552
553pub fn python_is_decorated_function_boundary(node: Node<'_>) -> bool {
554 if node.kind() != "decorated_definition" {
555 return false;
556 }
557 let mut cursor = node.walk();
558 node.named_children(&mut cursor)
559 .any(|child| child.kind() == "function_definition")
560}
561
562#[derive(Clone)]
563pub struct Scope {
564 kind: ScopeKind,
565 path: String,
566 fq: FqName,
573 code_unit: Option<CodeUnit>,
574 method_receiver: Option<String>,
575}
576
577#[derive(Clone, Copy, PartialEq, Eq)]
578enum ScopeKind {
579 Class,
580 Function,
581}
582
583pub struct PythonVisitor<'a> {
584 pub file: &'a ProjectFile,
585 pub source: &'a str,
586 pub package_name: &'a str,
587 module_fq: &'a FqName,
588 pub parsed: &'a mut ParsedFile,
589 pub module: Option<CodeUnit>,
590 pub overload_decorators: &'a PythonOverloadDecoratorBindings,
591}
592
593struct PythonContainer<'tree> {
594 node: Node<'tree>,
595 scope: Vec<Scope>,
596 module_control_depth: usize,
597}
598
599enum PythonWork<'tree> {
600 Container(PythonContainer<'tree>),
601 Statement {
602 node: Node<'tree>,
603 scope: Vec<Scope>,
604 module_control_depth: usize,
605 },
606}
607
608impl<'a> PythonVisitor<'a> {
609 pub fn visit_container(
610 &mut self,
611 node: Node<'_>,
612 scope: &[Scope],
613 module_control_depth: usize,
614 ) {
615 let mut stack = vec![PythonWork::Container(PythonContainer {
616 node,
617 scope: scope.to_vec(),
618 module_control_depth,
619 })];
620 while let Some(work) = stack.pop() {
621 match work {
622 PythonWork::Container(container) => {
623 let mut cursor = container.node.walk();
624 let children = container
625 .node
626 .named_children(&mut cursor)
627 .collect::<Vec<_>>();
628 for child in children.into_iter().rev() {
629 stack.push(PythonWork::Statement {
630 node: child,
631 scope: container.scope.clone(),
632 module_control_depth: container.module_control_depth,
633 });
634 }
635 }
636 PythonWork::Statement {
637 node,
638 scope,
639 module_control_depth,
640 } => self.visit_statement(node, &scope, module_control_depth, &mut stack),
641 }
642 }
643 }
644
645 fn visit_statement<'tree>(
646 &mut self,
647 node: Node<'tree>,
648 scope: &[Scope],
649 module_control_depth: usize,
650 stack: &mut Vec<PythonWork<'tree>>,
651 ) {
652 match node.kind() {
653 "decorated_definition" => {
654 if let Some(definition) = node.child_by_field_name("definition") {
655 self.visit_definition(
656 definition,
657 Some(node),
658 scope,
659 module_control_depth,
660 stack,
661 );
662 }
663 }
664 "class_definition" | "function_definition" => {
665 self.visit_definition(node, None, scope, module_control_depth, stack)
666 }
667 "expression_statement" => {
668 self.visit_expression_statement(node, scope, module_control_depth)
669 }
670 "import_statement" | "import_from_statement" => self.visit_import_statement(node),
671 "if_statement" | "try_statement" | "with_statement" | "for_statement"
672 | "while_statement" => {
673 let next_depth = if scope.is_empty() {
674 module_control_depth + 1
675 } else {
676 module_control_depth
677 };
678 stack.push(PythonWork::Container(PythonContainer {
679 node,
680 scope: scope.to_vec(),
681 module_control_depth: next_depth,
682 }));
683 }
684 "elif_clause" | "else_clause" | "except_clause" | "finally_clause" => {
685 stack.push(PythonWork::Container(PythonContainer {
686 node,
687 scope: scope.to_vec(),
688 module_control_depth,
689 }));
690 }
691 "block" | "module" => stack.push(PythonWork::Container(PythonContainer {
692 node,
693 scope: scope.to_vec(),
694 module_control_depth,
695 })),
696 _ => {}
697 }
698 }
699
700 fn visit_definition<'tree>(
701 &mut self,
702 definition: Node<'tree>,
703 wrapper: Option<Node<'tree>>,
704 scope: &[Scope],
705 module_control_depth: usize,
706 stack: &mut Vec<PythonWork<'tree>>,
707 ) {
708 match definition.kind() {
709 "class_definition" => self.visit_class_definition(
710 definition,
711 wrapper.unwrap_or(definition),
712 scope,
713 module_control_depth,
714 stack,
715 ),
716 "function_definition" => self.visit_function_definition(
717 definition,
718 wrapper.unwrap_or(definition),
719 scope,
720 module_control_depth,
721 stack,
722 ),
723 _ => {}
724 }
725 }
726
727 fn visit_class_definition<'tree>(
728 &mut self,
729 node: Node<'tree>,
730 range_node: Node<'tree>,
731 scope: &[Scope],
732 module_control_depth: usize,
733 stack: &mut Vec<PythonWork<'tree>>,
734 ) {
735 let Some(name_node) = node.child_by_field_name("name") else {
736 return;
737 };
738 let name = py_node_text(name_node, self.source).trim();
739 if name.is_empty() {
740 return;
741 }
742
743 let capture = !scope.is_empty() || module_control_depth <= 1;
744
745 let short_name = scope
746 .last()
747 .map(|parent| format!("{}${name}", parent.path))
748 .unwrap_or_else(|| name.to_string());
749 let fq = match scope.last() {
755 Some(parent) => parent
756 .fq
757 .clone()
758 .with_pushed(py_segment(name, SegmentKind::Nested)),
759 None => self
760 .module_fq
761 .clone()
762 .with_pushed(py_segment(name, SegmentKind::Type)),
763 };
764 let code_unit = CodeUnit::new_fq(
765 self.file.clone(),
766 CodeUnitType::Class,
767 self.package_name.to_string(),
768 short_name.clone(),
769 fq.clone(),
770 );
771 if capture {
772 self.parsed
773 .replace_code_unit(code_unit.clone(), range_node, self.source, None, None);
774 self.parsed.add_signature(
775 code_unit.clone(),
776 python_class_signature(range_node, self.source),
777 );
778 if let Some(module) = &self.module
779 && scope.is_empty()
780 {
781 self.parsed.add_child(module.clone(), code_unit.clone());
782 }
783 if let Some(parent) = scope.last()
784 && let Some(parent_cu) = &parent.code_unit
785 {
786 self.parsed.add_child(parent_cu.clone(), code_unit.clone());
787 }
788 self.parsed.set_raw_supertypes(
789 code_unit.clone(),
790 extract_python_supertypes(node, self.source),
791 );
792 }
793
794 let mut next_scope = scope.to_vec();
795 if capture {
796 next_scope.push(Scope {
797 kind: ScopeKind::Class,
798 path: short_name,
799 fq,
800 code_unit: Some(code_unit),
801 method_receiver: None,
802 });
803 }
804 if let Some(body) = node.child_by_field_name("body") {
805 stack.push(PythonWork::Container(PythonContainer {
806 node: body,
807 scope: next_scope,
808 module_control_depth,
809 }));
810 }
811 }
812
813 fn visit_function_definition<'tree>(
814 &mut self,
815 node: Node<'tree>,
816 range_node: Node<'tree>,
817 scope: &[Scope],
818 module_control_depth: usize,
819 stack: &mut Vec<PythonWork<'tree>>,
820 ) {
821 let Some(name_node) = node.child_by_field_name("name") else {
822 return;
823 };
824 let name = py_node_text(name_node, self.source).trim();
825 if name.is_empty() {
826 return;
827 }
828
829 let capture = !python_is_property_mutator(range_node, self.source)
830 && ((scope.is_empty() && module_control_depth <= 1)
831 || scope
832 .last()
833 .is_some_and(|parent| parent.kind == ScopeKind::Class));
834 let short_name = if let Some(parent) = scope.last() {
835 match parent.kind {
836 ScopeKind::Class => format!("{}.{}", parent.path, name),
837 ScopeKind::Function => format!("{}${name}", parent.path),
838 }
839 } else {
840 name.to_string()
841 };
842 let fq = if let Some(parent) = scope.last() {
847 match parent.kind {
848 ScopeKind::Class => parent
849 .fq
850 .clone()
851 .with_pushed(py_segment(name, SegmentKind::Member)),
852 ScopeKind::Function => parent
853 .fq
854 .clone()
855 .with_pushed(py_segment(name, SegmentKind::Nested)),
856 }
857 } else {
858 self.module_fq
859 .clone()
860 .with_pushed(py_segment(name, SegmentKind::Member))
861 };
862
863 if capture {
864 let code_unit_type = if python_function_has_decorator(node, self.source, "property") {
865 CodeUnitType::Field
866 } else {
867 CodeUnitType::Function
868 };
869 let signature = node
870 .child_by_field_name("parameters")
871 .map(|parameters| py_node_text(parameters, self.source).trim().to_string());
872 let code_unit = CodeUnit::with_signature_and_fq(
873 self.file.clone(),
874 code_unit_type,
875 self.package_name.to_string(),
876 short_name.clone(),
877 signature,
878 false,
879 fq.clone(),
880 );
881 self.parsed
882 .replace_code_unit(code_unit.clone(), range_node, self.source, None, None);
883 let signature = python_function_signature(range_node, self.source);
884 self.parsed.add_signature_with_metadata(
885 code_unit.clone(),
886 python_signature_metadata(signature, node, self.source).with_declaration_only(
887 self.overload_decorators
888 .decorates_as_overload(node, self.source),
889 ),
890 );
891 if let Some(module) = &self.module
892 && scope.is_empty()
893 {
894 self.parsed.add_child(module.clone(), code_unit.clone());
895 }
896 if let Some(parent) = scope.last()
897 && parent.kind == ScopeKind::Class
898 && let Some(parent_cu) = &parent.code_unit
899 {
900 self.parsed.add_child(parent_cu.clone(), code_unit.clone());
901 }
902 let scope_code_unit = Some(code_unit);
903 let mut next_scope = scope.to_vec();
904 next_scope.push(Scope {
905 kind: ScopeKind::Function,
906 path: short_name,
907 fq,
908 code_unit: scope_code_unit,
909 method_receiver: scope
910 .last()
911 .is_some_and(|parent| parent.kind == ScopeKind::Class)
912 .then(|| python_instance_method_receiver_name(node, self.source))
913 .flatten(),
914 });
915 if let Some(body) = node.child_by_field_name("body") {
916 stack.push(PythonWork::Container(PythonContainer {
917 node: body,
918 scope: next_scope,
919 module_control_depth,
920 }));
921 }
922 return;
923 }
924
925 let mut next_scope = scope.to_vec();
926 next_scope.push(Scope {
927 kind: ScopeKind::Function,
928 path: short_name,
929 fq,
930 code_unit: None,
931 method_receiver: None,
932 });
933 if let Some(body) = node.child_by_field_name("body") {
934 stack.push(PythonWork::Container(PythonContainer {
935 node: body,
936 scope: next_scope,
937 module_control_depth,
938 }));
939 }
940 }
941
942 fn visit_expression_statement(
943 &mut self,
944 node: Node<'_>,
945 scope: &[Scope],
946 module_control_depth: usize,
947 ) {
948 let Some(assignment) = node.named_child(0) else {
949 return;
950 };
951 if assignment.kind() != "assignment" {
952 return;
953 }
954 let Some(left) = assignment.child_by_field_name("left") else {
955 return;
956 };
957 self.visit_instance_attribute_assignment(left, scope);
958 let names = collect_assigned_names(left, self.source);
959 for name in names {
960 let (short_name, fq) = if let Some(parent) = scope.last() {
961 if parent.kind != ScopeKind::Class {
962 continue;
963 }
964 (
965 format!("{}.{}", parent.path, name),
966 parent
967 .fq
968 .clone()
969 .with_pushed(py_segment(&name, SegmentKind::Member)),
970 )
971 } else if module_control_depth <= 1 {
972 (
973 name.clone(),
974 self.module_fq
975 .clone()
976 .with_pushed(py_segment(&name, SegmentKind::Member)),
977 )
978 } else {
979 continue;
980 };
981 let code_unit = CodeUnit::new_fq(
982 self.file.clone(),
983 CodeUnitType::Field,
984 self.package_name.to_string(),
985 short_name,
986 fq,
987 );
988 if scope
989 .last()
990 .is_some_and(|parent| parent.kind == ScopeKind::Class)
991 {
992 self.parsed
996 .add_code_unit(code_unit.clone(), node, self.source, None, None);
997 } else {
998 self.parsed
999 .replace_code_unit(code_unit.clone(), node, self.source, None, None);
1000 }
1001 self.parsed.add_signature(
1002 code_unit.clone(),
1003 py_node_text(node, self.source).trim().to_string(),
1004 );
1005 if let Some(module) = &self.module
1006 && scope.is_empty()
1007 {
1008 self.parsed.add_child(module.clone(), code_unit.clone());
1009 }
1010 if let Some(parent) = scope.last()
1011 && parent.kind == ScopeKind::Class
1012 && let Some(parent_cu) = &parent.code_unit
1013 {
1014 self.parsed.add_child(parent_cu.clone(), code_unit);
1015 }
1016 }
1017 }
1018
1019 fn visit_instance_attribute_assignment(&mut self, left: Node<'_>, scope: &[Scope]) {
1020 let Some(function) = scope
1021 .last()
1022 .filter(|scope| scope.kind == ScopeKind::Function)
1023 else {
1024 return;
1025 };
1026 let Some(receiver) = function.method_receiver.as_deref() else {
1027 return;
1028 };
1029 let Some(parent) = scope
1030 .get(scope.len().saturating_sub(2))
1031 .filter(|scope| scope.kind == ScopeKind::Class)
1032 else {
1033 return;
1034 };
1035 let Some(parent_cu) = parent.code_unit.clone() else {
1036 return;
1037 };
1038 for (name, node) in collect_self_assigned_attributes(left, self.source, receiver) {
1039 let code_unit = CodeUnit::new_fq(
1040 self.file.clone(),
1041 CodeUnitType::Field,
1042 self.package_name.to_string(),
1043 format!("{}.{}", parent.path, name),
1044 parent
1045 .fq
1046 .clone()
1047 .with_pushed(py_segment(&name, SegmentKind::Member)),
1048 );
1049 if !self.parsed.contains_declaration(&code_unit) {
1050 self.parsed.replace_code_unit(
1051 code_unit.clone(),
1052 node,
1053 self.source,
1054 Some(parent_cu.clone()),
1055 Some(parent_cu.clone()),
1056 );
1057 }
1058 self.parsed.add_signature(
1059 code_unit.clone(),
1060 py_node_text(left, self.source).trim().to_string(),
1061 );
1062 }
1063 }
1064
1065 fn visit_import_statement(&mut self, node: Node<'_>) {
1066 for info in python_import_infos_from_node(node, self.source) {
1067 self.parsed.imports.push(info);
1068 }
1069 }
1070}
1071
1072pub fn parse_python_file(file: &ProjectFile, source: &str, tree: &Tree) -> ParsedFile {
1076 let module_components = python_module_components(file);
1077 let module_name = module_components.join(".");
1078 let module_fq = python_module_fq_from_components(&module_components);
1079 let mut parsed = ParsedFile::new(module_name.clone());
1080 let root = tree.root_node();
1081
1082 collect_python_identifiers(root, source, &mut parsed.type_identifiers);
1083
1084 let module_code_unit = module_code_unit_from_fq(file, &module_components, module_fq.clone());
1085 if let Some(module) = module_code_unit.clone() {
1086 parsed.add_code_unit(module, root, source, None, None);
1087 }
1088
1089 let overload_decorators = PythonOverloadDecoratorBindings::collect(root, source);
1090 let mut visitor = PythonVisitor {
1091 file,
1092 source,
1093 package_name: &module_name,
1094 module_fq: &module_fq,
1095 parsed: &mut parsed,
1096 module: module_code_unit,
1097 overload_decorators: &overload_decorators,
1098 };
1099 visitor.visit_container(root, &[], 0);
1100
1101 parsed
1102}
1103
1104pub fn py_node_text<'a>(node: Node<'_>, source: &'a str) -> &'a str {
1105 brokk_bifrost_core::analyzer::common::node_source_text(node, source)
1106}
1107
1108pub fn python_module_name(file: &ProjectFile) -> String {
1109 python_module_components(file).join(".")
1110}
1111
1112pub fn module_code_unit(file: &ProjectFile, module_fq: &str) -> Option<CodeUnit> {
1113 if module_fq.is_empty() {
1114 return None;
1115 }
1116 let components = python_module_components(file);
1117 debug_assert_eq!(
1118 module_fq,
1119 components.join("."),
1120 "module_code_unit must be built from the file's path-derived Python module name"
1121 );
1122 let structured_fq = python_module_fq_from_components(&components);
1123 module_code_unit_from_fq(file, &components, structured_fq)
1124}
1125
1126fn module_code_unit_from_fq(
1127 file: &ProjectFile,
1128 components: &[String],
1129 structured_fq: FqName,
1130) -> Option<CodeUnit> {
1131 let (short_name, package_components) = components.split_last()?;
1132 let package_name = package_components.join(".");
1133 Some(CodeUnit::new_fq(
1134 file.clone(),
1135 CodeUnitType::Module,
1136 package_name,
1137 short_name.clone(),
1138 structured_fq,
1139 ))
1140}
1141
1142fn python_class_signature(node: Node<'_>, source: &str) -> String {
1143 python_header_with_decorators(node, source)
1144}
1145
1146fn python_function_signature(node: Node<'_>, source: &str) -> String {
1147 let header = python_header_with_decorators(node, source);
1148 if let Some((head, tail)) = header.rsplit_once('\n') {
1149 format!("{head}\n{tail} ...")
1150 } else {
1151 format!("{header} ...")
1152 }
1153}
1154
1155fn python_signature_metadata(signature: String, node: Node<'_>, source: &str) -> SignatureMetadata {
1156 let Some(parameters_node) = node.child_by_field_name("parameters") else {
1157 return SignatureMetadata::new(signature, Vec::new())
1158 .with_dispatch_extensibility(DispatchExtensibility::Open);
1159 };
1160 let parameter_text = py_node_text(parameters_node, source).trim();
1161 let Some(parameters_start) = signature.find(parameter_text) else {
1162 return SignatureMetadata::new(signature, Vec::new())
1163 .with_dispatch_extensibility(DispatchExtensibility::Open);
1164 };
1165 let parameters_end = parameters_start + parameter_text.len();
1166 let mut search_start = parameters_start;
1167 let parameters = python_parameter_label_nodes(parameters_node)
1168 .into_iter()
1169 .filter_map(|label_node| {
1170 let label = py_node_text(label_node, source).trim();
1171 if label.is_empty() || search_start > parameters_end {
1172 return None;
1173 }
1174 let haystack = signature.get(search_start..parameters_end)?;
1175 let relative_start = haystack.find(label)?;
1176 let start_byte = search_start + relative_start;
1177 let end_byte = start_byte + label.len();
1178 search_start = end_byte;
1179 Some(ParameterMetadata::new(label, start_byte, end_byte))
1180 })
1181 .collect();
1182 SignatureMetadata::new(signature, parameters)
1183 .with_dispatch_extensibility(DispatchExtensibility::Open)
1184}
1185
1186fn python_parameter_label_nodes(parameters_node: Node<'_>) -> Vec<Node<'_>> {
1187 let mut labels = Vec::new();
1188 let mut cursor = parameters_node.walk();
1189 for child in parameters_node.named_children(&mut cursor) {
1190 if let Some(label_node) = python_parameter_label_node(child) {
1191 labels.push(label_node);
1192 }
1193 }
1194 labels
1195}
1196
1197pub fn python_parameter_label_node(node: Node<'_>) -> Option<Node<'_>> {
1205 match node.kind() {
1206 "identifier" => Some(node),
1207 "typed_parameter"
1208 | "typed_default_parameter"
1209 | "default_parameter"
1210 | "list_splat_pattern"
1211 | "dictionary_splat_pattern"
1212 | "keyword_separator" => node.child_by_field_name("name").or_else(|| {
1213 let mut cursor = node.walk();
1214 node.named_children(&mut cursor)
1215 .find_map(python_parameter_label_node)
1216 }),
1217 _ => None,
1218 }
1219}
1220
1221fn python_is_property_mutator(node: Node<'_>, source: &str) -> bool {
1222 python_header_with_decorators(node, source)
1223 .lines()
1224 .map(str::trim)
1225 .filter(|line| line.starts_with('@'))
1226 .any(|decorator| decorator.ends_with(".setter") || decorator.ends_with(".deleter"))
1227}
1228
1229pub fn python_expanded_comment_start(source: &str, start_byte: usize) -> usize {
1230 let line_starts = compute_line_starts(source);
1231 let line_index = find_line_index_for_offset(&line_starts, start_byte);
1232
1233 let mut comment_start = start_byte;
1234 for line_idx in (0..line_index).rev() {
1235 let line_start = line_starts[line_idx];
1236 let line_end = line_starts
1237 .get(line_idx + 1)
1238 .copied()
1239 .unwrap_or(source.len());
1240 let line = &source[line_start..line_end];
1241 let trimmed = line.trim_start();
1242
1243 if trimmed.trim().is_empty() {
1244 continue;
1245 }
1246
1247 if trimmed.starts_with('#') {
1248 comment_start = line_start;
1249 continue;
1250 }
1251
1252 break;
1253 }
1254
1255 comment_start
1256}
1257
1258fn python_header_with_decorators(node: Node<'_>, source: &str) -> String {
1259 let raw = py_node_text(node, source);
1260 let lines: Vec<_> = raw
1261 .lines()
1262 .map(str::trim_end)
1263 .filter(|line| !line.trim().is_empty())
1264 .collect();
1265 let mut relevant = Vec::new();
1266 for line in lines {
1267 let trimmed = line.trim_start();
1268 if trimmed.starts_with('@')
1269 || trimmed.starts_with("def ")
1270 || trimmed.starts_with("async def ")
1271 || trimmed.starts_with("class ")
1272 {
1273 relevant.push(trimmed.to_string());
1274 if trimmed.starts_with("def ")
1275 || trimmed.starts_with("async def ")
1276 || trimmed.starts_with("class ")
1277 {
1278 break;
1279 }
1280 }
1281 }
1282 relevant.join("\n")
1283}
1284
1285fn extract_python_supertypes(node: Node<'_>, source: &str) -> Vec<String> {
1286 let Some(superclasses) = node.child_by_field_name("superclasses") else {
1287 return Vec::new();
1288 };
1289 let mut result = Vec::new();
1290 let mut cursor = superclasses.walk();
1291 for child in superclasses.named_children(&mut cursor) {
1292 match child.kind() {
1293 "identifier" | "attribute" => {
1294 let text = py_node_text(child, source).trim();
1295 if !text.is_empty() {
1296 result.push(text.to_string());
1297 }
1298 }
1299 _ => {}
1300 }
1301 }
1302 result
1303}
1304
1305fn collect_assigned_names(node: Node<'_>, source: &str) -> Vec<String> {
1306 let mut names = Vec::new();
1307 walk_named_tree_preorder(node, true, |node| {
1308 match node.kind() {
1309 "attribute" | "subscript" => WalkControl::SkipChildren,
1313 "identifier" => {
1314 let text = py_node_text(node, source).trim();
1315 if !text.is_empty() {
1316 names.push(text.to_string());
1317 }
1318 WalkControl::Continue
1319 }
1320 _ => WalkControl::Continue,
1321 }
1322 });
1323 names
1324}
1325
1326fn collect_self_assigned_attributes<'tree>(
1327 node: Node<'tree>,
1328 source: &str,
1329 receiver_name: &str,
1330) -> Vec<(String, Node<'tree>)> {
1331 let mut attributes = Vec::new();
1332 collect_direct_self_assigned_attributes(node, source, receiver_name, &mut attributes);
1333 attributes
1334}
1335
1336fn collect_direct_self_assigned_attributes<'tree>(
1337 node: Node<'tree>,
1338 source: &str,
1339 receiver_name: &str,
1340 attributes: &mut Vec<(String, Node<'tree>)>,
1341) {
1342 match node.kind() {
1343 "attribute" => {
1344 let Some(object) = node.child_by_field_name("object") else {
1345 return;
1346 };
1347 if object.kind() != "identifier" || py_node_text(object, source).trim() != receiver_name
1348 {
1349 return;
1350 }
1351 let Some(attribute) = node.child_by_field_name("attribute") else {
1352 return;
1353 };
1354 let name = py_node_text(attribute, source).trim();
1355 if !name.is_empty() {
1356 attributes.push((name.to_string(), attribute));
1357 }
1358 }
1359 "pattern_list" | "tuple" | "list" | "parenthesized_expression" => {
1360 let mut cursor = node.walk();
1361 for child in node.named_children(&mut cursor) {
1362 collect_direct_self_assigned_attributes(child, source, receiver_name, attributes);
1363 }
1364 }
1365 _ => {}
1366 }
1367}
1368
1369fn python_instance_method_receiver_name(node: Node<'_>, source: &str) -> Option<String> {
1370 if python_function_has_decorator(node, source, "staticmethod")
1371 || python_function_has_decorator(node, source, "classmethod")
1372 {
1373 return None;
1374 }
1375 python_first_parameter_name(node, source)
1376}
1377
1378fn python_function_has_decorator(node: Node<'_>, source: &str, decorator_name: &str) -> bool {
1379 let Some(parent) = node.parent() else {
1380 return false;
1381 };
1382 if parent.kind() != "decorated_definition" {
1383 return false;
1384 }
1385 let mut cursor = parent.walk();
1386 parent
1387 .named_children(&mut cursor)
1388 .filter(|child| child.kind() == "decorator")
1389 .filter_map(|decorator| decorator.named_child(0))
1390 .filter_map(expression_name_node)
1391 .any(|name| py_node_text(name, source).trim() == decorator_name)
1392}
1393
1394fn python_first_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
1395 let parameters = node.child_by_field_name("parameters")?;
1396 let mut cursor = parameters.walk();
1397 parameters
1398 .named_children(&mut cursor)
1399 .find_map(|child| python_parameter_name(child, source))
1400}
1401
1402fn python_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
1403 match node.kind() {
1404 "identifier" => Some(py_node_text(node, source).trim().to_string()),
1405 "typed_parameter"
1406 | "default_parameter"
1407 | "list_splat_pattern"
1408 | "dictionary_splat_pattern" => node
1409 .child_by_field_name("name")
1410 .or_else(|| {
1411 let mut cursor = node.walk();
1412 node.named_children(&mut cursor)
1413 .find(|child| child.kind() == "identifier")
1414 })
1415 .and_then(|name| python_parameter_name(name, source)),
1416 _ => None,
1417 }
1418 .filter(|name| !name.is_empty())
1419}
1420
1421pub fn collect_python_identifiers(node: Node<'_>, source: &str, identifiers: &mut HashSet<String>) {
1422 walk_named_tree_preorder(node, true, |node| {
1423 if node.kind() == "identifier" {
1424 let text = py_node_text(node, source).trim();
1425 if !text.is_empty() {
1426 identifiers.insert(text.to_string());
1427 }
1428 }
1429 WalkControl::Continue
1430 });
1431}
1432
1433pub fn parse_python_tree(source: &str) -> Option<Tree> {
1434 let mut parser = Parser::new();
1435 parser
1436 .set_language(&tree_sitter_python::LANGUAGE.into())
1437 .expect("failed to load python parser");
1438 parser.parse(source, None)
1439}