1use brokk_bifrost_core::analyzer::common::node_source_text;
7use brokk_bifrost_core::analyzer::structural::adapter_helpers::{
8 attach_argument_role_with_derived_name, attach_role_with_derived_name, attach_terminal_callee,
9 field_name_in_parent, first_named_child, nearest_ancestor, node_range,
10};
11use brokk_bifrost_core::analyzer::structural::callable::CallSiteContext;
12use brokk_bifrost_core::analyzer::structural::edges::{
13 DEEP_REFERENCE_EDGE_SUPPORT, ReferenceEdgeSupport,
14};
15use brokk_bifrost_core::analyzer::structural::kinds::{NormalizedKind, Role};
16use brokk_bifrost_core::analyzer::structural::materialization::{
17 DeclarationMaterializationSupport, PYTHON_MATERIALIZATION_SUPPORT,
18};
19use brokk_bifrost_core::analyzer::structural::occurrences::{
20 Namespace, OccurrenceRole, OccurrenceRoleSupport, default_occurrence_namespace,
21};
22use brokk_bifrost_core::analyzer::structural::resolution::{
23 BindingActivation, BindingKind, DEEP_LEXICAL_ENVIRONMENT_SUPPORT, HoistingClass,
24 LexicalEnvironmentSupport,
25};
26use brokk_bifrost_core::analyzer::structural::routes::{
27 CuratedExportSurface, DEEP_IDENTITY_AXES, IdentityRouteSupport, RouteHopKind,
28};
29use brokk_bifrost_core::analyzer::structural::spec::{EmbeddedLeafFact, RoleSink, StructuralSpec};
30use brokk_bifrost_core::analyzer::{Language, Range};
31use brokk_bifrost_core::cancellation::CancellationToken;
32use brokk_bifrost_core::hash::HashSet;
33use tree_sitter::Node;
34
35use crate::syntax::{
36 expression_name_node, python_deferred_annotation_identifier_ranges,
37 python_keyword_argument_label, python_node_is_in_annotation,
38};
39
40#[derive(Debug, Default)]
41pub struct PythonStructuralSpec;
42
43pub static PYTHON_STRUCTURAL_SPEC: PythonStructuralSpec = PythonStructuralSpec;
44
45pub const PYTHON_KIND_TABLE: &[(&str, NormalizedKind)] = &[
49 ("call", NormalizedKind::Call),
50 ("attribute", NormalizedKind::FieldAccess),
51 ("function_definition", NormalizedKind::Function),
52 ("lambda", NormalizedKind::Lambda),
53 ("class_definition", NormalizedKind::Class),
54 ("assignment", NormalizedKind::Assignment),
55 ("import_statement", NormalizedKind::Import),
56 ("import_from_statement", NormalizedKind::Import),
57 ("identifier", NormalizedKind::Identifier),
58 ("string", NormalizedKind::StringLiteral),
59 ("concatenated_string", NormalizedKind::StringLiteral),
60 ("integer", NormalizedKind::NumericLiteral),
61 ("float", NormalizedKind::NumericLiteral),
62 ("true", NormalizedKind::BooleanLiteral),
63 ("false", NormalizedKind::BooleanLiteral),
64 ("none", NormalizedKind::NullLiteral),
65 ("return_statement", NormalizedKind::Return),
66 ("raise_statement", NormalizedKind::Throw),
67 ("except_clause", NormalizedKind::Catch),
68 ("if_statement", NormalizedKind::If),
69 ("for_statement", NormalizedKind::ForLoop),
70 ("list", NormalizedKind::CollectionLiteral),
71 ("set", NormalizedKind::CollectionLiteral),
72 ("dictionary", NormalizedKind::CollectionLiteral),
73 ("tuple", NormalizedKind::CollectionLiteral),
74 ("while_statement", NormalizedKind::WhileLoop),
75 ("block", NormalizedKind::Block),
80 ("decorator", NormalizedKind::Decorator),
81];
82
83fn attach_decorators(sink: &mut RoleSink<'_>, definition: Node<'_>) {
86 let Some(parent) = definition.parent() else {
87 return;
88 };
89 if parent.kind() != "decorated_definition" {
90 return;
91 }
92 for index in 0..parent.named_child_count() {
93 let Some(child) = parent.named_child(index) else {
94 continue;
95 };
96 if child.kind() == "decorator" {
97 attach_role_with_derived_name(sink, Role::Decorator, child, expression_name_node);
98 }
99 }
100}
101
102static PYTHON_OCCURRENCE_ROLE_SUPPORT: OccurrenceRoleSupport = OccurrenceRoleSupport::NONE
103 .supported(OccurrenceRole::DeclarationName)
104 .supported(OccurrenceRole::Binder)
105 .supported(OccurrenceRole::LabelOrKey)
106 .supported(OccurrenceRole::TypeOperand)
107 .supported(OccurrenceRole::PathSegment)
108 .supported(OccurrenceRole::ImportAlias)
109 .supported(OccurrenceRole::ImportTarget)
110 .supported(OccurrenceRole::ReceiverPosition)
111 .supported(OccurrenceRole::MemberPosition)
112 .supported(OccurrenceRole::ValueReference);
113
114fn python_dotted_name_is_import(dotted_name: Node<'_>) -> bool {
121 let mut current = dotted_name;
122 loop {
123 let Some(parent) = current.parent() else {
124 return false;
125 };
126 match parent.kind() {
127 "import_statement" | "import_from_statement" | "future_import_statement" => {
128 return true;
129 }
130 "aliased_import" | "dotted_name" | "relative_import" => current = parent,
131 _ => return false,
132 }
133 }
134}
135
136fn python_occurrence_role(node: Node<'_>) -> Option<OccurrenceRole> {
142 if node.kind() != "identifier" {
143 return None;
144 }
145 let parent = node.parent()?;
146 let field = field_name_in_parent(parent, node);
147 let role = match parent.kind() {
148 "function_definition" | "class_definition" if field == Some("name") => {
149 OccurrenceRole::DeclarationName
150 }
151 "type" | "generic_type" | "type_parameter" | "constrained_type" | "union_type" => {
154 OccurrenceRole::TypeOperand
155 }
156 "parameters"
157 | "lambda_parameters"
158 | "typed_parameter"
159 | "list_splat_pattern"
160 | "dictionary_splat_pattern"
161 | "tuple_pattern"
162 | "list_pattern"
163 | "pattern_list"
164 | "as_pattern_target" => OccurrenceRole::Binder,
165 "default_parameter" | "typed_default_parameter" if field == Some("name") => {
166 OccurrenceRole::Binder
167 }
168 "for_statement" | "for_in_clause" if field == Some("left") => OccurrenceRole::Binder,
169 "keyword_argument" if python_keyword_argument_label(node) => OccurrenceRole::LabelOrKey,
170 "attribute" => match field {
171 Some("attribute") => OccurrenceRole::MemberPosition,
172 Some("object") => OccurrenceRole::ReceiverPosition,
173 _ => OccurrenceRole::ValueReference,
174 },
175 "aliased_import" if field == Some("alias") => OccurrenceRole::ImportAlias,
176 "import_from_statement" if field == Some("name") => OccurrenceRole::ImportTarget,
177 "dotted_name" => {
178 let is_tail =
179 parent.named_child(parent.named_child_count().saturating_sub(1)) == Some(node);
180 match (is_tail, python_dotted_name_is_import(parent)) {
181 (true, true) => OccurrenceRole::ImportTarget,
182 (true, false) => OccurrenceRole::ValueReference,
183 (false, _) => OccurrenceRole::PathSegment,
184 }
185 }
186 _ => OccurrenceRole::ValueReference,
187 };
188 Some(role)
189}
190
191fn python_definition_is_method(definition: Node<'_>) -> bool {
201 let mut current = definition;
202 while let Some(parent) = current.parent() {
203 match parent.kind() {
204 "class_definition" => return true,
205 "block" | "decorated_definition" => current = parent,
208 _ => return false,
209 }
210 }
211 false
212}
213
214fn python_binding_activation(binder: Node<'_>, scope: Range) -> Option<BindingActivation> {
225 let Some(form) = nearest_ancestor(binder, |kind| {
233 matches!(
234 kind,
235 "parameters"
236 | "lambda_parameters"
237 | "for_statement"
238 | "for_in_clause"
239 | "as_pattern"
240 | "list_comprehension"
241 | "set_comprehension"
242 | "dictionary_comprehension"
243 | "generator_expression"
244 | "function_definition"
245 )
246 }) else {
247 return Some(BindingActivation {
248 kind: BindingKind::Local,
249 hoisting: HoistingClass::ScopeWide,
250 activation: scope,
251 });
252 };
253 match form.kind() {
254 "parameters" | "lambda_parameters" | "function_definition" => Some(BindingActivation {
255 kind: BindingKind::Parameter,
256 hoisting: HoistingClass::ScopeWide,
257 activation: scope,
258 }),
259 "for_in_clause" => {
260 let comprehension = nearest_ancestor(form, |kind| {
262 matches!(
263 kind,
264 "list_comprehension"
265 | "set_comprehension"
266 | "dictionary_comprehension"
267 | "generator_expression"
268 )
269 })?;
270 Some(BindingActivation {
271 kind: BindingKind::LoopVariable,
272 hoisting: HoistingClass::DeclaredHead,
273 activation: node_range(comprehension),
274 })
275 }
276 "for_statement" => Some(BindingActivation {
277 kind: BindingKind::LoopVariable,
278 hoisting: HoistingClass::ScopeWide,
279 activation: scope,
280 }),
281 "as_pattern" => Some(BindingActivation {
282 kind: BindingKind::PatternBinder,
283 hoisting: HoistingClass::ScopeWide,
284 activation: scope,
285 }),
286 _ => Some(BindingActivation {
287 kind: BindingKind::Local,
288 hoisting: HoistingClass::ScopeWide,
289 activation: scope,
290 }),
291 }
292}
293
294fn python_plain_string_text<'a>(node: Node<'_>, source: &'a str) -> Option<&'a str> {
299 if node.kind() != "string" {
300 return None;
301 }
302 let mut cursor = node.walk();
303 let mut content = None;
304 for child in node.named_children(&mut cursor) {
305 match child.kind() {
306 "string_start" | "string_end" => {}
307 "string_content" if content.is_none() && child.named_child_count() == 0 => {
308 content = Some(child);
309 }
310 _ => return None,
311 }
312 }
313 Some(content.map_or("", |child| node_source_text(child, source)))
315}
316
317#[derive(Debug, Clone, Copy, PartialEq, Eq)]
319enum ReadableSurface {
320 Yes,
321 No,
322}
323
324fn python_collect_all_members(
328 value: Node<'_>,
329 source: &str,
330 names: &mut HashSet<String>,
331) -> ReadableSurface {
332 if !matches!(value.kind(), "list" | "tuple") {
333 return ReadableSurface::No;
334 }
335 let mut cursor = value.walk();
336 for element in value.named_children(&mut cursor) {
337 match python_plain_string_text(element, source) {
338 Some(text) => {
339 names.insert(text.to_owned());
340 }
341 None => return ReadableSurface::No,
342 }
343 }
344 ReadableSurface::Yes
345}
346
347fn python_curated_export_surface(root: Node<'_>, source: &str) -> CuratedExportSurface {
357 let names_all = |node: Option<Node<'_>>| {
358 node.is_some_and(|node| {
359 node.kind() == "identifier" && node_source_text(node, source) == "__all__"
360 })
361 };
362 let mut names: HashSet<String> = HashSet::default();
363 let mut stated = false;
364 let mut readable = ReadableSurface::Yes;
365 let mut cursor = root.walk();
366 for statement in root.named_children(&mut cursor) {
367 if statement.kind() != "expression_statement" {
368 continue;
369 }
370 let mut inner = statement.walk();
371 for expression in statement.named_children(&mut inner) {
372 match expression.kind() {
373 "assignment" | "augmented_assignment" => {
374 if !names_all(expression.child_by_field_name("left")) {
375 continue;
376 }
377 stated = true;
378 let Some(value) = expression.child_by_field_name("right") else {
381 continue;
382 };
383 let extends = expression.kind() == "assignment"
386 || expression
387 .child_by_field_name("operator")
388 .is_some_and(|operator| operator.kind() == "+=");
389 if !extends
390 || python_collect_all_members(value, source, &mut names)
391 == ReadableSurface::No
392 {
393 readable = ReadableSurface::No;
394 }
395 }
396 "call" => {
399 let Some(function) = expression.child_by_field_name("function") else {
400 continue;
401 };
402 if function.kind() == "attribute"
403 && names_all(function.child_by_field_name("object"))
404 {
405 stated = true;
406 readable = ReadableSurface::No;
407 }
408 }
409 _ => {}
410 }
411 }
412 }
413 match (stated, readable) {
414 (false, _) => CuratedExportSurface::Absent,
415 (true, ReadableSurface::Yes) => CuratedExportSurface::Listed(names),
416 (true, ReadableSurface::No) => CuratedExportSurface::Unreadable,
417 }
418}
419
420fn python_indirection_relation(
445 token: Node<'_>,
446 source: &str,
447 surface: &CuratedExportSurface,
448) -> Option<RouteHopKind> {
449 let statement = nearest_ancestor(token, |kind| {
450 matches!(
451 kind,
452 "import_statement" | "import_from_statement" | "future_import_statement"
453 )
454 })?;
455 let mut clause = token;
458 while let Some(parent) = clause.parent() {
459 if parent.id() == statement.id() {
460 break;
461 }
462 clause = parent;
463 }
464
465 if field_name_in_parent(statement, clause) == Some("module_name") {
466 let mut cursor = statement.walk();
470 let star = statement
471 .children(&mut cursor)
472 .any(|child| child.kind() == "wildcard_import");
473 return Some(if star {
474 RouteHopKind::ReExport
475 } else {
476 RouteHopKind::Import
477 });
478 }
479
480 let bound = match clause.kind() {
481 "aliased_import" => {
482 let name = clause.child_by_field_name("name")?;
483 let alias = clause.child_by_field_name("alias")?;
484 if node_source_text(name, source) == node_source_text(alias, source) {
485 return Some(RouteHopKind::ReExport);
486 }
487 alias
488 }
489 "dotted_name" if statement.kind() == "import_statement" => clause.named_child(0)?,
492 "dotted_name" => clause,
493 _ => return None,
494 };
495 match surface.lists(node_source_text(bound, source)) {
496 Some(true) => Some(RouteHopKind::ReExport),
497 Some(false) => Some(RouteHopKind::Import),
498 None => None,
499 }
500}
501
502impl StructuralSpec for PythonStructuralSpec {
503 fn language(&self) -> Language {
504 Language::Python
505 }
506
507 fn supports_boolean_literal_value(&self) -> bool {
508 true
509 }
510
511 fn reference_edge_support(&self) -> &ReferenceEdgeSupport {
512 &DEEP_REFERENCE_EDGE_SUPPORT
513 }
514
515 fn identity_route_support(&self) -> &IdentityRouteSupport {
516 static SUPPORT: IdentityRouteSupport = DEEP_IDENTITY_AXES
519 .supported_relation(RouteHopKind::Alias)
520 .supported_relation(RouteHopKind::Import)
521 .supported_relation(RouteHopKind::ReExport)
522 .supported_relation(RouteHopKind::NestedOwner);
523 &SUPPORT
524 }
525
526 fn qualified_path_root<'tree>(&self, token: Node<'tree>) -> Option<Node<'tree>> {
529 if token.kind() != "identifier" {
530 return None;
531 }
532 token
533 .parent()
534 .filter(|parent| parent.kind() == "dotted_name")
535 }
536
537 fn path_segment_tokens<'tree>(&self, root: Node<'tree>) -> Vec<Node<'tree>> {
538 if root.kind() != "dotted_name" {
539 return Vec::new();
540 }
541 let mut cursor = root.walk();
542 root.named_children(&mut cursor)
543 .filter(|child| child.kind() == "identifier")
544 .collect()
545 }
546
547 fn curated_export_surface(&self, root: Node<'_>, source: &str) -> CuratedExportSurface {
548 python_curated_export_surface(root, source)
549 }
550
551 fn indirection_relation(
552 &self,
553 token: Node<'_>,
554 source: &str,
555 surface: &CuratedExportSurface,
556 ) -> Option<RouteHopKind> {
557 python_indirection_relation(token, source, surface)
558 }
559
560 fn kind_table(&self) -> &'static [(&'static str, NormalizedKind)] {
561 PYTHON_KIND_TABLE
562 }
563
564 fn refine_kind(
565 &self,
566 node: Node<'_>,
567 kind: NormalizedKind,
568 _enclosing: Option<NormalizedKind>,
569 _source: &str,
570 _context: &CallSiteContext,
571 ) -> NormalizedKind {
572 if kind == NormalizedKind::Function && python_definition_is_method(node) {
573 NormalizedKind::Method
574 } else {
575 kind
576 }
577 }
578
579 fn should_extract(&self, node: Node<'_>, kind: NormalizedKind) -> bool {
580 kind != NormalizedKind::Assignment || node.child_by_field_name("right").is_some()
581 }
582
583 fn supports_kind(&self, kind: NormalizedKind) -> bool {
584 kind == NormalizedKind::Method
585 || self
586 .kind_table()
587 .iter()
588 .any(|(_, fact_kind)| fact_kind.satisfies(kind))
589 }
590
591 fn occurrence_role_support(&self) -> &OccurrenceRoleSupport {
592 &PYTHON_OCCURRENCE_ROLE_SUPPORT
593 }
594
595 fn lexical_environment_support(&self) -> &LexicalEnvironmentSupport {
596 &DEEP_LEXICAL_ENVIRONMENT_SUPPORT
597 }
598
599 fn materialization_support(&self) -> &DeclarationMaterializationSupport {
600 &PYTHON_MATERIALIZATION_SUPPORT
601 }
602
603 fn binding_activation(&self, binder: Node<'_>, scope: Range) -> Option<BindingActivation> {
604 python_binding_activation(binder, scope)
605 }
606
607 fn occurrence_namespace(
610 &self,
611 role: OccurrenceRole,
612 declares: Option<NormalizedKind>,
613 ) -> Option<Namespace> {
614 match role {
615 OccurrenceRole::PathSegment => Some(Namespace::Module),
616 _ => default_occurrence_namespace(role, declares),
617 }
618 }
619
620 fn embedded_leaf_facts(
621 &self,
622 node: Node<'_>,
623 kind: NormalizedKind,
624 source: &str,
625 cancellation: Option<&CancellationToken>,
626 ) -> Vec<EmbeddedLeafFact> {
627 if kind != NormalizedKind::StringLiteral
628 || node.kind() != "string"
629 || !python_node_is_in_annotation(node)
630 {
631 return Vec::new();
632 }
633
634 python_deferred_annotation_identifier_ranges(node, source, cancellation)
635 .unwrap_or_default()
636 .into_iter()
637 .map(|range| EmbeddedLeafFact {
638 kind: NormalizedKind::Identifier,
639 range,
640 occurrence_role: OccurrenceRole::TypeOperand,
641 })
642 .collect()
643 }
644
645 fn extract(&self, node: Node<'_>, kind: NormalizedKind, sink: &mut RoleSink<'_>) {
646 if let Some(role) = python_occurrence_role(node) {
647 sink.occurrence_role(node, role);
648 }
649 match kind {
650 NormalizedKind::Call => {
651 if let Some(function) = node.child_by_field_name("function") {
652 attach_terminal_callee(sink, function, expression_name_node(function));
655 if function.kind() == "attribute"
656 && let Some(object) = function.child_by_field_name("object")
657 {
658 attach_role_with_derived_name(
659 sink,
660 Role::Receiver,
661 object,
662 expression_name_node,
663 );
664 }
665 }
666 if let Some(arguments) = node.child_by_field_name("arguments") {
667 for index in 0..arguments.named_child_count() {
668 if !sink.should_continue() {
669 break;
670 }
671 let Some(argument) = arguments.named_child(index) else {
672 continue;
673 };
674 match argument.kind() {
675 "comment" => {}
676 "keyword_argument" => {
677 if let (Some(keyword), Some(value)) = (
678 argument.child_by_field_name("name"),
679 argument.child_by_field_name("value"),
680 ) {
681 sink.kwarg(keyword, value);
682 }
683 }
684 _ => attach_argument_role_with_derived_name(
685 sink,
686 argument,
687 expression_name_node,
688 ),
689 }
690 }
691 }
692 }
693 NormalizedKind::FieldAccess => {
694 if let Some(attribute) = node.child_by_field_name("attribute") {
695 sink.set_name(attribute);
696 sink.role_named(Role::Field, attribute, attribute);
697 }
698 if let Some(object) = node.child_by_field_name("object") {
699 attach_role_with_derived_name(sink, Role::Object, object, expression_name_node);
700 }
701 }
702 NormalizedKind::Function | NormalizedKind::Method | NormalizedKind::Class => {
703 if let Some(name) = node.child_by_field_name("name") {
704 sink.set_name(name);
705 }
706 attach_decorators(sink, node);
707 }
708 NormalizedKind::Assignment => {
709 if let Some(left) = node.child_by_field_name("left") {
710 attach_role_with_derived_name(sink, Role::Left, left, expression_name_node);
711 }
712 if let Some(right) = node.child_by_field_name("right") {
713 attach_role_with_derived_name(sink, Role::Right, right, expression_name_node);
714 }
715 }
716 NormalizedKind::Import => match node.kind() {
717 "import_from_statement" => {
718 if let Some(module) = node.child_by_field_name("module_name") {
719 sink.role_named(Role::Module, module, module);
720 }
721 }
722 _ => {
723 for index in 0..node.named_child_count() {
724 if !sink.should_continue() {
725 break;
726 }
727 let Some(child) = node.named_child(index) else {
728 continue;
729 };
730 match child.kind() {
731 "dotted_name" => sink.role_named(Role::Module, child, child),
732 "aliased_import" => {
733 if let Some(name) = child.child_by_field_name("name") {
734 sink.role_named(Role::Module, name, name);
735 }
736 }
737 _ => {}
738 }
739 }
740 }
741 },
742 NormalizedKind::Identifier => sink.set_name(node),
743 NormalizedKind::Decorator => {
744 if let Some(name) = first_named_child(node).and_then(expression_name_node) {
745 sink.set_name(name);
746 }
747 }
748 NormalizedKind::ForLoop => {
749 if let Some(right) = node.child_by_field_name("right") {
750 attach_role_with_derived_name(
751 sink,
752 Role::Iterable,
753 right,
754 expression_name_node,
755 );
756 }
757 }
758 NormalizedKind::CollectionLiteral => {
759 for index in 0..node.named_child_count() {
760 let Some(child) = node.named_child(index) else {
761 continue;
762 };
763 if child.kind() == "comment" {
764 continue;
765 }
766 attach_role_with_derived_name(sink, Role::Element, child, expression_name_node);
767 }
768 }
769 _ => {}
770 }
771 }
772}