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, python_plain_string_literal,
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
294#[derive(Debug, Clone, Copy, PartialEq, Eq)]
300enum ReadableSurface {
301 Yes,
302 No,
303}
304
305fn python_collect_all_members(
309 value: Node<'_>,
310 source: &str,
311 names: &mut HashSet<String>,
312) -> ReadableSurface {
313 if !matches!(value.kind(), "list" | "tuple") {
314 return ReadableSurface::No;
315 }
316 let mut cursor = value.walk();
317 for element in value.named_children(&mut cursor) {
318 match python_plain_string_literal(element, source) {
319 Some(text) => {
320 names.insert(text.to_owned());
321 }
322 None => return ReadableSurface::No,
323 }
324 }
325 ReadableSurface::Yes
326}
327
328fn python_curated_export_surface(root: Node<'_>, source: &str) -> CuratedExportSurface {
338 let names_all = |node: Option<Node<'_>>| {
339 node.is_some_and(|node| {
340 node.kind() == "identifier" && node_source_text(node, source) == "__all__"
341 })
342 };
343 let mut names: HashSet<String> = HashSet::default();
344 let mut stated = false;
345 let mut readable = ReadableSurface::Yes;
346 let mut cursor = root.walk();
347 for statement in root.named_children(&mut cursor) {
348 if statement.kind() != "expression_statement" {
349 continue;
350 }
351 let mut inner = statement.walk();
352 for expression in statement.named_children(&mut inner) {
353 match expression.kind() {
354 "assignment" | "augmented_assignment" => {
355 if !names_all(expression.child_by_field_name("left")) {
356 continue;
357 }
358 stated = true;
359 let Some(value) = expression.child_by_field_name("right") else {
362 continue;
363 };
364 let extends = expression.kind() == "assignment"
367 || expression
368 .child_by_field_name("operator")
369 .is_some_and(|operator| operator.kind() == "+=");
370 if !extends
371 || python_collect_all_members(value, source, &mut names)
372 == ReadableSurface::No
373 {
374 readable = ReadableSurface::No;
375 }
376 }
377 "call" => {
380 let Some(function) = expression.child_by_field_name("function") else {
381 continue;
382 };
383 if function.kind() == "attribute"
384 && names_all(function.child_by_field_name("object"))
385 {
386 stated = true;
387 readable = ReadableSurface::No;
388 }
389 }
390 _ => {}
391 }
392 }
393 }
394 match (stated, readable) {
395 (false, _) => CuratedExportSurface::Absent,
396 (true, ReadableSurface::Yes) => CuratedExportSurface::Listed(names),
397 (true, ReadableSurface::No) => CuratedExportSurface::Unreadable,
398 }
399}
400
401fn python_indirection_relation(
426 token: Node<'_>,
427 source: &str,
428 surface: &CuratedExportSurface,
429) -> Option<RouteHopKind> {
430 let statement = nearest_ancestor(token, |kind| {
431 matches!(
432 kind,
433 "import_statement" | "import_from_statement" | "future_import_statement"
434 )
435 })?;
436 let mut clause = token;
439 while let Some(parent) = clause.parent() {
440 if parent.id() == statement.id() {
441 break;
442 }
443 clause = parent;
444 }
445
446 if field_name_in_parent(statement, clause) == Some("module_name") {
447 let mut cursor = statement.walk();
451 let star = statement
452 .children(&mut cursor)
453 .any(|child| child.kind() == "wildcard_import");
454 return Some(if star {
455 RouteHopKind::ReExport
456 } else {
457 RouteHopKind::Import
458 });
459 }
460
461 let bound = match clause.kind() {
462 "aliased_import" => {
463 let name = clause.child_by_field_name("name")?;
464 let alias = clause.child_by_field_name("alias")?;
465 if node_source_text(name, source) == node_source_text(alias, source) {
466 return Some(RouteHopKind::ReExport);
467 }
468 alias
469 }
470 "dotted_name" if statement.kind() == "import_statement" => clause.named_child(0)?,
473 "dotted_name" => clause,
474 _ => return None,
475 };
476 match surface.lists(node_source_text(bound, source)) {
477 Some(true) => Some(RouteHopKind::ReExport),
478 Some(false) => Some(RouteHopKind::Import),
479 None => None,
480 }
481}
482
483impl StructuralSpec for PythonStructuralSpec {
484 fn language(&self) -> Language {
485 Language::Python
486 }
487
488 fn supports_boolean_literal_value(&self) -> bool {
489 true
490 }
491
492 fn reference_edge_support(&self) -> &ReferenceEdgeSupport {
493 &DEEP_REFERENCE_EDGE_SUPPORT
494 }
495
496 fn identity_route_support(&self) -> &IdentityRouteSupport {
497 static SUPPORT: IdentityRouteSupport = DEEP_IDENTITY_AXES
500 .supported_relation(RouteHopKind::Alias)
501 .supported_relation(RouteHopKind::Import)
502 .supported_relation(RouteHopKind::ReExport)
503 .supported_relation(RouteHopKind::NestedOwner);
504 &SUPPORT
505 }
506
507 fn qualified_path_root<'tree>(&self, token: Node<'tree>) -> Option<Node<'tree>> {
510 if token.kind() != "identifier" {
511 return None;
512 }
513 token
514 .parent()
515 .filter(|parent| parent.kind() == "dotted_name")
516 }
517
518 fn path_segment_tokens<'tree>(&self, root: Node<'tree>) -> Vec<Node<'tree>> {
519 if root.kind() != "dotted_name" {
520 return Vec::new();
521 }
522 let mut cursor = root.walk();
523 root.named_children(&mut cursor)
524 .filter(|child| child.kind() == "identifier")
525 .collect()
526 }
527
528 fn curated_export_surface(&self, root: Node<'_>, source: &str) -> CuratedExportSurface {
529 python_curated_export_surface(root, source)
530 }
531
532 fn indirection_relation(
533 &self,
534 token: Node<'_>,
535 source: &str,
536 surface: &CuratedExportSurface,
537 ) -> Option<RouteHopKind> {
538 python_indirection_relation(token, source, surface)
539 }
540
541 fn kind_table(&self) -> &'static [(&'static str, NormalizedKind)] {
542 PYTHON_KIND_TABLE
543 }
544
545 fn refine_kind(
546 &self,
547 node: Node<'_>,
548 kind: NormalizedKind,
549 _enclosing: Option<NormalizedKind>,
550 _source: &str,
551 _context: &CallSiteContext,
552 ) -> NormalizedKind {
553 if kind == NormalizedKind::Function && python_definition_is_method(node) {
554 NormalizedKind::Method
555 } else {
556 kind
557 }
558 }
559
560 fn should_extract(&self, node: Node<'_>, kind: NormalizedKind) -> bool {
561 kind != NormalizedKind::Assignment || node.child_by_field_name("right").is_some()
562 }
563
564 fn supports_kind(&self, kind: NormalizedKind) -> bool {
565 kind == NormalizedKind::Method
566 || self
567 .kind_table()
568 .iter()
569 .any(|(_, fact_kind)| fact_kind.satisfies(kind))
570 }
571
572 fn occurrence_role_support(&self) -> &OccurrenceRoleSupport {
573 &PYTHON_OCCURRENCE_ROLE_SUPPORT
574 }
575
576 fn lexical_environment_support(&self) -> &LexicalEnvironmentSupport {
577 &DEEP_LEXICAL_ENVIRONMENT_SUPPORT
578 }
579
580 fn materialization_support(&self) -> &DeclarationMaterializationSupport {
581 &PYTHON_MATERIALIZATION_SUPPORT
582 }
583
584 fn binding_activation(&self, binder: Node<'_>, scope: Range) -> Option<BindingActivation> {
585 python_binding_activation(binder, scope)
586 }
587
588 fn occurrence_namespace(
591 &self,
592 role: OccurrenceRole,
593 declares: Option<NormalizedKind>,
594 ) -> Option<Namespace> {
595 match role {
596 OccurrenceRole::PathSegment => Some(Namespace::Module),
597 _ => default_occurrence_namespace(role, declares),
598 }
599 }
600
601 fn embedded_leaf_facts(
602 &self,
603 node: Node<'_>,
604 kind: NormalizedKind,
605 source: &str,
606 cancellation: Option<&CancellationToken>,
607 ) -> Vec<EmbeddedLeafFact> {
608 if kind != NormalizedKind::StringLiteral
609 || node.kind() != "string"
610 || !python_node_is_in_annotation(node)
611 {
612 return Vec::new();
613 }
614
615 python_deferred_annotation_identifier_ranges(node, source, cancellation)
616 .unwrap_or_default()
617 .into_iter()
618 .map(|range| EmbeddedLeafFact {
619 kind: NormalizedKind::Identifier,
620 range,
621 occurrence_role: OccurrenceRole::TypeOperand,
622 })
623 .collect()
624 }
625
626 fn extract(&self, node: Node<'_>, kind: NormalizedKind, sink: &mut RoleSink<'_>) {
627 if let Some(role) = python_occurrence_role(node) {
628 sink.occurrence_role(node, role);
629 }
630 match kind {
631 NormalizedKind::Call => {
632 if let Some(function) = node.child_by_field_name("function") {
633 attach_terminal_callee(sink, function, expression_name_node(function));
636 if function.kind() == "attribute"
637 && let Some(object) = function.child_by_field_name("object")
638 {
639 attach_role_with_derived_name(
640 sink,
641 Role::Receiver,
642 object,
643 expression_name_node,
644 );
645 }
646 }
647 if let Some(arguments) = node.child_by_field_name("arguments") {
648 for index in 0..arguments.named_child_count() {
649 if !sink.should_continue() {
650 break;
651 }
652 let Some(argument) = arguments.named_child(index) else {
653 continue;
654 };
655 match argument.kind() {
656 "comment" => {}
657 "keyword_argument" => {
658 if let (Some(keyword), Some(value)) = (
659 argument.child_by_field_name("name"),
660 argument.child_by_field_name("value"),
661 ) {
662 sink.kwarg(keyword, value);
663 }
664 }
665 _ => attach_argument_role_with_derived_name(
666 sink,
667 argument,
668 expression_name_node,
669 ),
670 }
671 }
672 }
673 }
674 NormalizedKind::FieldAccess => {
675 if let Some(attribute) = node.child_by_field_name("attribute") {
676 sink.set_name(attribute);
677 sink.role_named(Role::Field, attribute, attribute);
678 }
679 if let Some(object) = node.child_by_field_name("object") {
680 attach_role_with_derived_name(sink, Role::Object, object, expression_name_node);
681 }
682 }
683 NormalizedKind::Function | NormalizedKind::Method | NormalizedKind::Class => {
684 if let Some(name) = node.child_by_field_name("name") {
685 sink.set_name(name);
686 }
687 attach_decorators(sink, node);
688 }
689 NormalizedKind::Assignment => {
690 if let Some(left) = node.child_by_field_name("left") {
691 attach_role_with_derived_name(sink, Role::Left, left, expression_name_node);
692 }
693 if let Some(right) = node.child_by_field_name("right") {
694 attach_role_with_derived_name(sink, Role::Right, right, expression_name_node);
695 }
696 }
697 NormalizedKind::Import => match node.kind() {
698 "import_from_statement" => {
699 if let Some(module) = node.child_by_field_name("module_name") {
700 sink.role_named(Role::Module, module, module);
701 }
702 }
703 _ => {
704 for index in 0..node.named_child_count() {
705 if !sink.should_continue() {
706 break;
707 }
708 let Some(child) = node.named_child(index) else {
709 continue;
710 };
711 match child.kind() {
712 "dotted_name" => sink.role_named(Role::Module, child, child),
713 "aliased_import" => {
714 if let Some(name) = child.child_by_field_name("name") {
715 sink.role_named(Role::Module, name, name);
716 }
717 }
718 _ => {}
719 }
720 }
721 }
722 },
723 NormalizedKind::Identifier => sink.set_name(node),
724 NormalizedKind::Decorator => {
725 if let Some(name) = first_named_child(node).and_then(expression_name_node) {
726 sink.set_name(name);
727 }
728 }
729 NormalizedKind::ForLoop => {
730 if let Some(right) = node.child_by_field_name("right") {
731 attach_role_with_derived_name(
732 sink,
733 Role::Iterable,
734 right,
735 expression_name_node,
736 );
737 }
738 }
739 NormalizedKind::CollectionLiteral => {
740 for index in 0..node.named_child_count() {
741 let Some(child) = node.named_child(index) else {
742 continue;
743 };
744 if child.kind() == "comment" {
745 continue;
746 }
747 attach_role_with_derived_name(sink, Role::Element, child, expression_name_node);
748 }
749 }
750 _ => {}
751 }
752 }
753}