1use brokk_bifrost_core::analyzer::Range;
2use brokk_bifrost_core::analyzer::common::node_source_text as node_text;
3use brokk_bifrost_core::cancellation::CancellationToken;
4use brokk_bifrost_core::hash::HashSet;
5use tree_sitter::{Node, Tree};
6
7#[cfg(any(test, feature = "test-support"))]
10pub(crate) fn is_python_member_path(value: &str) -> bool {
11 static IDENTIFIER: std::sync::LazyLock<::regex::Regex> = std::sync::LazyLock::new(|| {
12 ::regex::Regex::new(r"\A[_\p{XID_Start}]\p{XID_Continue}*\z")
13 .expect("Python identifier value regex")
14 });
15 value.contains('.') && value.split('.').all(|part| IDENTIFIER.is_match(part))
16}
17
18pub(crate) fn python_conditional_expression_operands(node: Node<'_>) -> Option<[Node<'_>; 3]> {
26 assert_eq!(node.kind(), "conditional_expression");
27 let mut cursor = node.walk();
28 let operands: Vec<Node<'_>> = node
29 .named_children(&mut cursor)
30 .filter(|child| !child.is_extra())
31 .collect();
32 operands.try_into().ok()
33}
34
35pub(crate) fn python_parameter_name(node: Node<'_>, source: &str) -> Option<String> {
36 python_parameter_name_bounded(node, source, || true)
37 .expect("unbounded parameter name lookup completes")
38}
39
40pub(crate) fn python_parameter_name_bounded(
41 node: Node<'_>,
42 source: &str,
43 scope_step: impl FnMut() -> bool,
44) -> Option<Option<String>> {
45 python_parameter_node_bounded(node, scope_step).map(|node| {
46 node.and_then(|node| {
47 let name = node_text(node, source).trim();
48 (!name.is_empty()).then(|| name.to_owned())
49 })
50 })
51}
52
53#[cfg(any(test, feature = "test-support"))]
55pub(crate) fn python_parameter_node(node: Node<'_>) -> Option<Node<'_>> {
56 python_parameter_node_bounded(node, || true).expect("unbounded parameter token lookup")
57}
58
59fn python_parameter_node_bounded(
60 mut node: Node<'_>,
61 mut scope_step: impl FnMut() -> bool,
62) -> Option<Option<Node<'_>>> {
63 loop {
64 if !scope_step() {
65 return None;
66 }
67 match node.kind() {
68 "identifier" => return Some((!node.is_missing()).then_some(node)),
69 "typed_parameter"
70 | "typed_default_parameter"
71 | "default_parameter"
72 | "list_splat_pattern"
73 | "dictionary_splat_pattern" => {
74 if let Some(name) = node.child_by_field_name("name") {
75 node = name;
76 continue;
77 }
78 let mut selected = None;
79 let mut cursor = node.walk();
80 for child in node.named_children(&mut cursor) {
81 if !scope_step() {
82 return None;
83 }
84 if matches!(
85 child.kind(),
86 "identifier" | "list_splat_pattern" | "dictionary_splat_pattern"
87 ) {
88 selected = Some(child);
89 break;
90 }
91 }
92 let Some(name) = selected else {
93 return Some(None);
94 };
95 node = name;
96 }
97 _ => return Some(None),
98 }
99 }
100}
101
102pub fn leftmost_identifier(mut node: Node<'_>) -> Option<Node<'_>> {
103 loop {
104 match node.kind() {
105 "identifier" => return Some(node),
106 "attribute" => node = node.child_by_field_name("object")?,
107 _ => return None,
108 }
109 }
110}
111
112pub(crate) fn attribute_chain<'tree>(node: Node<'tree>) -> Option<(Node<'tree>, Vec<Node<'tree>>)> {
113 let mut attributes = Vec::new();
114 let mut current = node;
115 loop {
116 if current.kind() != "attribute" {
117 return None;
118 }
119 attributes.push(current.child_by_field_name("attribute")?);
120 current = current.child_by_field_name("object")?;
121 if current.kind() == "identifier" {
122 attributes.reverse();
123 return Some((current, attributes));
124 }
125 }
126}
127
128pub fn python_reference_is_deferred_function_body(node: Node<'_>) -> bool {
129 python_reference_is_deferred_function_body_bounded(node, || true)
130 .expect("unbounded deferred-context lookup completes")
131}
132
133pub fn python_reference_is_deferred_function_body_bounded(
135 node: Node<'_>,
136 mut scope_step: impl FnMut() -> bool,
137) -> Option<bool> {
138 let site_start = node.start_byte();
139 let site_end = node.end_byte();
140 let mut current = node;
141 while let Some(parent) = current.parent() {
142 if !scope_step() {
143 return None;
144 }
145 if matches!(parent.kind(), "function_definition" | "lambda")
146 && parent
147 .child_by_field_name("body")
148 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
149 {
150 return Some(true);
151 }
152 current = parent;
153 }
154 Some(false)
155}
156
157pub fn python_static_attribute_path<'tree>(node: Node<'tree>) -> Option<Vec<Node<'tree>>> {
163 python_static_value_path_with(node, || true, |_| None)
164 .expect("unbounded static path lookup completes")
165}
166
167pub fn python_static_attribute_path_bounded<'tree>(
170 node: Node<'tree>,
171 scope_step: impl FnMut() -> bool,
172) -> Option<Vec<Node<'tree>>> {
173 python_static_value_path_with(node, scope_step, |_| None).flatten()
174}
175
176pub fn python_static_value_path_bounded<'tree>(
179 node: Node<'tree>,
180 scope_step: impl FnMut() -> bool,
181) -> Option<Option<Vec<Node<'tree>>>> {
182 python_static_value_path_with(
183 node,
184 scope_step,
185 crate::structural::python_parenthesized_expression_child,
186 )
187}
188
189fn python_static_value_path_with<'tree>(
190 mut node: Node<'tree>,
191 mut scope_step: impl FnMut() -> bool,
192 mut parenthesized_child: impl FnMut(Node<'tree>) -> Option<Node<'tree>>,
193) -> Option<Option<Vec<Node<'tree>>>> {
194 let mut path = Vec::new();
195 loop {
196 if !scope_step() {
197 return None;
198 }
199 match node.kind() {
200 "identifier" => {
201 path.push(node);
202 break;
203 }
204 "attribute" => {
205 let Some(object) = push_static_attribute_segment(node, &mut path) else {
206 return Some(None);
207 };
208 node = object;
209 }
210 "parenthesized_expression" => {
211 for _ in brokk_bifrost_core::analyzer::tree_walk::named_children_iter(node) {
214 if !scope_step() {
215 return None;
216 }
217 }
218 let Some(child) = parenthesized_child(node) else {
219 return Some(None);
220 };
221 node = child;
222 }
223 _ => return Some(None),
224 }
225 }
226 for _ in 0..path.len() {
227 if !scope_step() {
228 return None;
229 }
230 }
231 path.reverse();
232 Some(Some(path))
233}
234
235pub fn python_static_type_path<'tree>(mut node: Node<'tree>) -> Option<Vec<Node<'tree>>> {
239 let mut path = Vec::new();
240 loop {
241 match node.kind() {
242 "identifier" => {
243 path.push(node);
244 break;
245 }
246 "type" | "generic_type" | "subscript" => node = node.named_child(0)?,
247 "attribute" => {
248 node = push_static_attribute_segment(node, &mut path)?;
249 }
250 "member_type" => {
251 let mut cursor = node.walk();
252 let mut children = node.named_children(&mut cursor);
253 let qualifier = children.next()?;
254 let member = children.next()?;
255 if member.kind() != "identifier" || children.next().is_some() {
256 return None;
257 }
258 path.push(member);
259 node = qualifier;
260 }
261 _ => return None,
262 }
263 }
264 path.reverse();
265 Some(path)
266}
267
268fn push_static_attribute_segment<'tree>(
269 node: Node<'tree>,
270 path: &mut Vec<Node<'tree>>,
271) -> Option<Node<'tree>> {
272 let attribute = node.child_by_field_name("attribute")?;
273 if attribute.kind() != "identifier" {
274 return None;
275 }
276 path.push(attribute);
277 node.child_by_field_name("object")
278}
279
280pub fn python_plain_string_literal<'source>(
286 node: Node<'_>,
287 source: &'source str,
288) -> Option<&'source str> {
289 if node.kind() != "string"
290 || node
291 .parent()
292 .is_some_and(|parent| parent.kind() == "concatenated_string")
293 {
294 return None;
295 }
296 let mut content = None;
297 let mut cursor = node.walk();
298 for child in node.named_children(&mut cursor) {
299 match child.kind() {
300 "string_start" | "string_end" => {
301 let delimiter = child.utf8_text(source.as_bytes()).ok()?;
302 if delimiter
303 .chars()
304 .any(|character| character != '"' && character != '\'')
305 {
306 return None;
307 }
308 }
309 "string_content" if child.named_child_count() == 0 && content.is_none() => {
310 content = Some(child);
311 }
312 _ => return None,
313 }
314 }
315 Some(content.map_or("", |child| {
316 child
317 .utf8_text(source.as_bytes())
318 .expect("a tree-sitter node range is valid UTF-8 source")
319 }))
320}
321
322#[derive(Debug, Default)]
323pub struct PythonOverloadDecoratorBindings {
324 direct: HashSet<String>,
325 namespaces: HashSet<String>,
326}
327
328#[derive(Debug)]
329pub(crate) enum PythonOverloadDecoratorName {
330 Direct(String),
331 Namespace(String),
332}
333
334impl PythonOverloadDecoratorBindings {
335 pub(crate) fn collect_import(
338 &mut self,
339 import: &brokk_bifrost_core::analyzer::model::ImportInfo,
340 ) {
341 use brokk_bifrost_core::analyzer::model::StructuredImportPathKind;
342 let Some(path) = &import.path else { return };
343 match (path.kind, path.segments.as_slice()) {
344 (Some(StructuredImportPathKind::Namespace), [module]) if is_typing_module(module) => {
345 self.namespaces
346 .insert(import.alias.as_ref().unwrap_or(module).clone());
347 }
348 (Some(StructuredImportPathKind::ImportFrom), [module, name])
349 if is_typing_module(module) && name == "overload" =>
350 {
351 self.direct
352 .insert(import.alias.as_ref().unwrap_or(name).clone());
353 }
354 _ => {}
355 }
356 }
357
358 pub(crate) fn overload_decorator_names(
363 function: Node<'_>,
364 source: &str,
365 ) -> Vec<PythonOverloadDecoratorName> {
366 let Some(parent) = function
367 .parent()
368 .filter(|node| node.kind() == "decorated_definition")
369 else {
370 return Vec::new();
371 };
372 let mut cursor = parent.walk();
373 parent
374 .named_children(&mut cursor)
375 .filter(|child| child.kind() == "decorator")
376 .filter_map(decorator_callee)
377 .filter_map(|callee| match callee.kind() {
378 "identifier" => Some(PythonOverloadDecoratorName::Direct(
379 node_text(callee, source).trim().to_string(),
380 )),
381 "attribute" => {
382 let attribute = callee.child_by_field_name("attribute")?;
383 if node_text(attribute, source).trim() != "overload" {
384 return None;
385 }
386 let object = callee.child_by_field_name("object")?;
387 (object.kind() == "identifier").then(|| {
388 PythonOverloadDecoratorName::Namespace(
389 node_text(object, source).trim().to_string(),
390 )
391 })
392 }
393 _ => None,
394 })
395 .collect()
396 }
397
398 pub(crate) fn matches_overload_binding(&self, name: &PythonOverloadDecoratorName) -> bool {
399 match name {
400 PythonOverloadDecoratorName::Direct(name) => self.direct.contains(name),
401 PythonOverloadDecoratorName::Namespace(name) => self.namespaces.contains(name),
402 }
403 }
404}
405
406fn is_typing_module(module: &str) -> bool {
407 matches!(module, "typing" | "typing_extensions")
408}
409
410pub fn expression_name_node<'tree>(expression: Node<'tree>) -> Option<Node<'tree>> {
412 let mut current = expression;
413 loop {
414 match current.kind() {
415 "identifier" => return Some(current),
416 "attribute" => current = current.child_by_field_name("attribute")?,
417 "call" => current = current.child_by_field_name("function")?,
418 _ => return None,
419 }
420 }
421}
422
423pub fn python_keyword_argument_label(node: Node<'_>) -> bool {
431 node.kind() == "identifier"
432 && node.parent().is_some_and(|parent| {
433 parent.kind() == "keyword_argument" && parent.child_by_field_name("name") == Some(node)
434 })
435}
436
437pub fn decorator_callee<'tree>(decorator: Node<'tree>) -> Option<Node<'tree>> {
439 if decorator.kind() != "decorator" {
440 return None;
441 }
442 let mut cursor = decorator.walk();
445 let mut expression = decorator
446 .named_children(&mut cursor)
447 .find(|child| !child.is_extra())?;
448 while expression.kind() == "call" {
449 expression = expression.child_by_field_name("function")?;
450 }
451 Some(expression)
452}
453
454pub fn python_node_is_in_annotation(node: Node<'_>) -> bool {
457 let start = node.start_byte();
458 let end = node.end_byte();
459 let mut current = node;
460 while let Some(parent) = current.parent() {
461 let annotation = match parent.kind() {
462 "function_definition" => parent.child_by_field_name("return_type"),
463 "typed_parameter" | "typed_default_parameter" | "assignment" => {
464 parent.child_by_field_name("type")
465 }
466 _ => None,
467 };
468 if let Some(annotation) = annotation
469 && annotation.start_byte() <= start
470 && end <= annotation.end_byte()
471 {
472 return true;
473 }
474 current = parent;
475 }
476 false
477}
478
479pub fn python_deferred_annotation_identifier_ranges(
481 string: Node<'_>,
482 source: &str,
483 cancellation: Option<&CancellationToken>,
484) -> Option<Vec<Range>> {
485 let tree = python_deferred_annotation_tree(string, source, cancellation)?;
486
487 let mut ranges = Vec::new();
488 let mut stack = vec![tree.root_node()];
489 while let Some(current) = stack.pop() {
490 if cancellation.is_some_and(CancellationToken::is_cancelled) {
491 return None;
492 }
493 if current.kind() == "identifier" {
494 ranges.push(Range {
495 start_byte: current.start_byte(),
496 end_byte: current.end_byte(),
497 start_line: current.start_position().row + 1,
498 end_line: current.end_position().row + 1,
499 });
500 }
501 for index in (0..current.named_child_count()).rev() {
502 if let Some(child) = current.named_child(index) {
503 stack.push(child);
504 }
505 }
506 }
507 Some(ranges)
508}
509
510pub fn python_deferred_annotation_tree(
514 string: Node<'_>,
515 source: &str,
516 cancellation: Option<&CancellationToken>,
517) -> Option<Tree> {
518 if string.kind() != "string"
519 || string
520 .parent()
521 .is_some_and(|parent| parent.kind() == "concatenated_string")
522 || !python_node_is_in_annotation(string)
523 || python_string_is_literal_value(string, source)
524 {
525 return None;
526 }
527
528 let mut content = None;
529 for index in 0..string.named_child_count() {
530 let child = string.named_child(index)?;
531 match child.kind() {
532 "string_start" | "string_end" => {}
533 "string_content" if content.is_none() => content = Some(child),
534 _ => return None,
535 }
536 }
537 let content = content?;
538 let language = tree_sitter_python::LANGUAGE.into();
539 let tree = brokk_bifrost_core::analyzer::common::parse_source_range_with_cancellation(
540 &language,
541 source,
542 content.range(),
543 cancellation,
544 )?;
545 if tree.root_node().has_error() {
546 return None;
547 }
548 Some(tree)
549}
550
551fn python_string_is_literal_value(string: Node<'_>, source: &str) -> bool {
555 let start = string.start_byte();
556 let end = string.end_byte();
557 let mut current = string;
558 while let Some(parent) = current.parent() {
559 match parent.kind() {
560 "subscript" => {
561 let Some(value) = parent.child_by_field_name("value") else {
562 return false;
563 };
564 if value.start_byte() <= start && end <= value.end_byte() {
565 return false;
566 }
567 return python_literal_annotation_base(value, source);
568 }
569 "generic_type" => {
570 let Some(value) = parent.named_child(0) else {
571 return false;
572 };
573 return python_literal_annotation_base(value, source);
574 }
575 _ => current = parent,
576 }
577 }
578 false
579}
580
581fn python_literal_annotation_base(value: Node<'_>, source: &str) -> bool {
582 match value.kind() {
583 "identifier" => node_text(value, source) == "Literal",
584 "attribute" => {
585 let (Some(object), Some(attribute)) = (
586 value.child_by_field_name("object"),
587 value.child_by_field_name("attribute"),
588 ) else {
589 return false;
590 };
591 object.kind() == "identifier"
592 && matches!(node_text(object, source), "typing" | "typing_extensions")
593 && attribute.kind() == "identifier"
594 && node_text(attribute, source) == "Literal"
595 }
596 "member_type" => {
597 let mut identifiers = Vec::new();
598 let mut stack = vec![value];
599 while let Some(node) = stack.pop() {
600 if node.kind() == "identifier" {
601 identifiers.push(node_text(node, source));
602 continue;
603 }
604 for index in (0..node.named_child_count()).rev() {
605 if let Some(child) = node.named_child(index) {
606 stack.push(child);
607 }
608 }
609 }
610 matches!(
611 identifiers.as_slice(),
612 ["typing", "Literal"] | ["typing_extensions", "Literal"]
613 )
614 }
615 _ => false,
616 }
617}
618
619pub(crate) fn python_contains_all_identifier(node: Node<'_>, source: &str) -> bool {
622 use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
623 let mut found = false;
624 walk_named_tree_preorder(node, true, |node| {
625 if node.kind() == "identifier" && &source[node.byte_range()] == "__all__" {
626 found = true;
627 WalkControl::Break
628 } else {
629 WalkControl::Continue
630 }
631 });
632 found
633}
634
635#[cfg(test)]
636mod tests {
637 use brokk_bifrost_test_macros::hardwon;
638
639 #[test]
640 fn qs6_python_syntax_field_boundary() {
641 use super::*;
642 for (source, expected) in [
643 ("sys.argv[1]", Some("sys.argv")),
644 ("obj.member[key]", Some("obj.member")),
645 ("obj.nested.member[key]", None),
646 ("values[key]", None),
647 ("obj.member", None),
648 ] {
649 let tree = crate::declarations::parse_python_tree(source).unwrap();
650 assert!(!tree.root_node().has_error());
651 let expression = tree
652 .root_node()
653 .named_child(0)
654 .unwrap()
655 .named_child(0)
656 .unwrap();
657 let seed = crate::runtime_values::python_runtime_keyed_access_seed(expression);
658 assert_eq!(seed.map(|node| &source[node.byte_range()]), expected);
659 if let Some(seed) = seed {
660 assert_eq!(seed.kind(), "attribute");
661 assert_eq!(
662 expression.child_by_field_name("value").unwrap().id(),
663 seed.id()
664 );
665 assert_eq!(seed.child_by_field_name("object").unwrap().start_byte(), 0);
666 assert_eq!(seed.start_byte(), 0);
667 assert_eq!(seed.end_byte(), expected.unwrap().len());
668 }
669 }
670 for (annotation, names) in [
671 ("pkg.Type", vec!["pkg", "Type"]),
672 ("pkg.Outer.Inner[int]", vec!["pkg", "Outer", "Inner"]),
673 ("Outer[T].Inner", vec!["Outer", "Inner"]),
674 ] {
675 let source = format!("def f() -> {annotation}: pass\n");
676 let tree = crate::declarations::parse_python_tree(&source).unwrap();
677 assert!(!tree.root_node().has_error());
678 let annotation_node = tree
679 .root_node()
680 .named_child(0)
681 .unwrap()
682 .child_by_field_name("return_type")
683 .unwrap();
684 let path = python_static_type_path(annotation_node).unwrap();
685 assert_eq!(
686 path.iter()
687 .map(|node| &source[node.byte_range()])
688 .collect::<Vec<_>>(),
689 names
690 );
691 for node in path {
692 assert_eq!(node.kind(), "identifier");
693 assert!(node.start_byte() >= annotation_node.start_byte());
694 assert!(node.end_byte() <= annotation_node.end_byte());
695 }
696 }
697 let source = "@factory()()\ndef f(): pass\n";
698 let tree = crate::declarations::parse_python_tree(source).unwrap();
699 let decorator = tree
700 .root_node()
701 .named_child(0)
702 .unwrap()
703 .named_child(0)
704 .unwrap();
705 assert_eq!(decorator.kind(), "decorator");
706 let call = decorator.named_child(0).unwrap();
707 assert_eq!(call.kind(), "call");
708 assert_eq!(call.child_by_field_name("function").unwrap().kind(), "call");
709 let callee = decorator_callee(decorator).unwrap();
710 assert_eq!(callee.kind(), "identifier");
711 assert_eq!(callee.byte_range(), 1..8);
712 assert_eq!(&source[callee.byte_range()], "factory");
713 }
714
715 #[test]
716 fn qs6_deferred_annotation_live_token_ranges() {
717 use super::*;
718 let source = "# \u{e9}\ndef f(x: \"pkg.Type[Other]\") -> \"Result\": pass\nordinary = \"text\"\ny: Literal[\"value\"]\n";
719 let tree = crate::declarations::parse_python_tree(source).unwrap();
720 assert!(!tree.root_node().has_error());
721 let fresh = CancellationToken::new();
722 let cancelled = CancellationToken::new();
723 cancelled.cancel();
724 let mut pending = vec![tree.root_node()];
725 let mut strings = Vec::new();
726 while let Some(node) = pending.pop() {
727 if node.kind() == "string" {
728 strings.push(node);
729 }
730 for index in (0..node.named_child_count()).rev() {
731 pending.push(node.named_child(index).unwrap());
732 }
733 }
734 assert_eq!(strings.len(), 4);
735 let mut names = Vec::new();
736 for string in strings {
737 let no_token = python_deferred_annotation_identifier_ranges(string, source, None);
738 assert_eq!(
739 python_deferred_annotation_identifier_ranges(string, source, Some(&fresh)),
740 no_token
741 );
742 assert!(
743 python_deferred_annotation_identifier_ranges(string, source, Some(&cancelled))
744 .is_none()
745 );
746 if let Some(ranges) = no_token {
747 for range in ranges {
748 assert!(range.start_byte > string.start_byte());
749 assert!(range.end_byte < string.end_byte());
750 assert_eq!(range.start_line, 2);
751 assert_eq!(range.end_line, 2);
752 names.push(&source[range.start_byte..range.end_byte]);
753 }
754 }
755 }
756 assert_eq!(names, ["pkg", "Type", "Other", "Result"]);
757 }
758
759 #[test]
760 #[hardwon("#5044: commented parentheses must preserve import identity and budget exhaustion")]
761 fn static_paths_keep_parentheses_comments_and_exhaustion_distinct() {
762 let source = "( # receiver\n module).function\n";
763 let tree = crate::declarations::parse_python_tree(source).unwrap();
764 let expression = tree
765 .root_node()
766 .named_child(0)
767 .unwrap()
768 .named_child(0)
769 .unwrap();
770 assert!(super::python_static_attribute_path(expression).is_none());
772 assert!(super::python_static_value_path_bounded(expression, || false).is_none());
773 let path = super::python_static_value_path_bounded(expression, || true)
774 .unwrap()
775 .unwrap();
776 assert_eq!(
777 path.iter()
778 .map(|node| node.utf8_text(source.as_bytes()).unwrap())
779 .collect::<Vec<_>>(),
780 ["module", "function"]
781 );
782 let dynamic = crate::declarations::parse_python_tree("factory().function\n").unwrap();
783 let expression = dynamic
784 .root_node()
785 .named_child(0)
786 .unwrap()
787 .named_child(0)
788 .unwrap();
789 assert!(matches!(
790 super::python_static_value_path_bounded(expression, || true),
791 Some(None)
792 ));
793 }
794}