1use brokk_bifrost_core::analyzer::Range;
14use tree_sitter::Node;
15
16use crate::bindings::{
17 PythonDirectScopeBindingKind, PythonLexicalNameResolution,
18 python_direct_scope_bindings_bounded, python_name_resolution_at,
19};
20use crate::imports::python_import_bindings_from_tree;
21use crate::syntax::python_plain_string_literal;
22
23const PROCESS_INPUT_CONTAINERS: [(&str, &str); 2] = [("os", "environ"), ("sys", "argv")];
31
32#[derive(Clone, Debug, Eq, PartialEq)]
37pub enum PythonRuntimeAccessKey {
38 Property(String),
40 Index(u128),
42 Dynamic,
44 Unsupported,
47}
48
49#[derive(Clone, Copy, Debug, Eq, PartialEq)]
51pub enum PythonRuntimeRootResolution {
52 ImportedModule,
55 LexicallyBound,
58 Unknown,
60}
61
62#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum PythonRuntimeMutationEvidence {
65 NoKnownWrite,
66 KnownWrite,
67 KnownWriteAndUnknownEffects,
68 UnknownEffects,
69}
70
71#[derive(Clone, Copy, Debug, Eq, PartialEq)]
74pub enum PythonRuntimeAccessorCoverage {
75 NoKnownAccessorEffects,
76 UnknownAccessorEffects,
77}
78
79#[derive(Clone, Copy, Debug, Eq, PartialEq)]
81pub enum PythonRuntimeExecutionContext {
82 Module,
83 DeferredFunction(Range),
84}
85
86#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88pub enum PythonRuntimeWriteTarget {
89 Container,
91 Key,
93 ContainerEscape,
96 ModuleEscape,
99}
100
101#[derive(Clone, Debug, Eq, PartialEq)]
103pub struct PythonRuntimeWrite {
104 pub root_name: String,
106 pub container: Option<String>,
109 pub range: Range,
110 pub target: PythonRuntimeWriteTarget,
111}
112
113#[derive(Clone, Debug, Eq, PartialEq)]
115pub struct PythonRuntimeRead {
116 pub binding_name: String,
118 pub root_name: String,
120 pub container: String,
121 pub access: PythonRuntimeAccessKey,
122 pub range: Range,
124 pub root_range: Range,
125 pub container_range: Range,
129 pub candidate_anchor: Range,
130 pub key_range: Range,
131 pub lexical_resolution: PythonRuntimeRootResolution,
132 pub mutation: PythonRuntimeMutationEvidence,
133 pub execution: PythonRuntimeExecutionContext,
134 pub accessor: PythonRuntimeAccessorCoverage,
138}
139
140#[derive(Clone, Debug, Eq, PartialEq)]
142pub struct PythonRuntimeReadFacts {
143 pub reads: Vec<PythonRuntimeRead>,
144 pub writes: Vec<PythonRuntimeWrite>,
145 pub visited_nodes: usize,
147 pub complete: bool,
150}
151
152#[derive(Debug)]
154struct PythonRuntimeBudget {
155 remaining: usize,
156 exhausted: bool,
157 visited_nodes: usize,
158}
159
160impl PythonRuntimeBudget {
161 fn for_facts(max_facts: usize) -> Self {
164 Self {
165 remaining: max_facts.saturating_mul(64).max(4096),
166 exhausted: false,
167 visited_nodes: 0,
168 }
169 }
170
171 fn visit(&mut self) -> bool {
172 if self.remaining == 0 {
173 self.exhausted = true;
174 return false;
175 }
176 self.remaining -= 1;
177 self.visited_nodes += 1;
178 true
179 }
180}
181
182struct ModuleBinding {
184 local_name: String,
185 root_module: String,
187 scope_start_byte: usize,
188 scope_end_byte: usize,
189 function_scoped: bool,
190}
191
192pub fn extract_python_runtime_reads(
198 root: Node<'_>,
199 source: &str,
200 max_facts: usize,
201) -> PythonRuntimeReadFacts {
202 let mut budget = PythonRuntimeBudget::for_facts(max_facts);
203 let bindings = module_bindings(root, source);
204 let mut reads = Vec::new();
205 let mut writes = Vec::new();
206 let mut stack = vec![root];
207 while let Some(node) = stack.pop() {
208 if !budget.visit() {
209 break;
210 }
211 if node.kind() == "attribute"
212 && reads.len() + writes.len() < max_facts
213 && let Some(path) = reviewed_attribute_path(node, source)
214 {
215 let resolution = root_resolution(
216 path.root,
217 &path.root_module,
218 &bindings,
219 root,
220 source,
221 &mut budget,
222 );
223 if resolution == PythonRuntimeRootResolution::ImportedModule {
224 classify_container_occurrence(
225 node,
226 &path,
227 source,
228 root,
229 &mut budget,
230 &mut reads,
231 &mut writes,
232 );
233 }
234 }
235 if node.kind() == "identifier"
236 && writes.len() + reads.len() < max_facts
237 && let Some(module) =
238 escaping_module_identifier(node, &bindings, root, source, &mut budget)
239 {
240 writes.push(PythonRuntimeWrite {
241 root_name: module,
242 container: None,
243 range: node_source_range(node),
244 target: PythonRuntimeWriteTarget::ModuleEscape,
245 });
246 }
247 let mut cursor = node.walk();
248 stack.extend(node.named_children(&mut cursor));
249 }
250 for read in &mut reads {
253 let known_write = writes.iter().any(|write| {
254 write.root_name == read.root_name
255 && write
256 .container
257 .as_ref()
258 .is_none_or(|container| *container == read.container)
259 });
260 read.mutation = match (known_write, read.accessor) {
261 (false, PythonRuntimeAccessorCoverage::NoKnownAccessorEffects) => {
262 PythonRuntimeMutationEvidence::NoKnownWrite
263 }
264 (false, PythonRuntimeAccessorCoverage::UnknownAccessorEffects) => {
265 PythonRuntimeMutationEvidence::UnknownEffects
266 }
267 (true, PythonRuntimeAccessorCoverage::NoKnownAccessorEffects) => {
268 PythonRuntimeMutationEvidence::KnownWrite
269 }
270 (true, PythonRuntimeAccessorCoverage::UnknownAccessorEffects) => {
271 PythonRuntimeMutationEvidence::KnownWriteAndUnknownEffects
272 }
273 };
274 }
275 reads.sort_by_key(|read| read.range.start_byte);
276 writes.sort_by_key(|write| write.range.start_byte);
277 PythonRuntimeReadFacts {
278 complete: !budget.exhausted && reads.len() + writes.len() < max_facts,
279 visited_nodes: budget.visited_nodes,
280 reads,
281 writes,
282 }
283}
284
285pub fn python_runtime_keyed_access_seed<'tree>(node: Node<'tree>) -> Option<Node<'tree>> {
293 if node.kind() != "subscript" {
294 return None;
295 }
296 let base = node.child_by_field_name("value")?;
297 if base.kind() != "attribute" {
298 return None;
299 }
300 let attribute = base.child_by_field_name("attribute")?;
301 let object = base.child_by_field_name("object")?;
302 if object.kind() != "identifier" || attribute.kind() != "identifier" {
303 return None;
304 }
305 Some(base)
306}
307
308struct ReviewedPath<'tree> {
310 root: Node<'tree>,
311 root_module: String,
312 container: String,
313}
314
315fn reviewed_attribute_path<'tree>(node: Node<'tree>, source: &str) -> Option<ReviewedPath<'tree>> {
318 let object = node.child_by_field_name("object")?;
319 let attribute = node.child_by_field_name("attribute")?;
320 if object.kind() != "identifier" || attribute.kind() != "identifier" {
321 return None;
322 }
323 let container = node_text(attribute, source);
324 let (root_module, _) = PROCESS_INPUT_CONTAINERS
325 .iter()
326 .find(|(_, reviewed)| *reviewed == container)?;
327 Some(ReviewedPath {
328 root: object,
329 root_module: (*root_module).to_owned(),
330 container: container.to_owned(),
331 })
332}
333
334fn classify_container_occurrence<'tree>(
337 attribute: Node<'tree>,
338 path: &ReviewedPath<'tree>,
339 source: &str,
340 module_root: Node<'tree>,
341 budget: &mut PythonRuntimeBudget,
342 reads: &mut Vec<PythonRuntimeRead>,
343 writes: &mut Vec<PythonRuntimeWrite>,
344) {
345 let write = |target| PythonRuntimeWrite {
346 root_name: path.root_module.clone(),
347 container: Some(path.container.clone()),
348 range: node_source_range(attribute),
349 target,
350 };
351 let Some(parent) = attribute.parent() else {
352 writes.push(write(PythonRuntimeWriteTarget::ContainerEscape));
353 return;
354 };
355 if parent.kind() != "subscript"
356 || parent
357 .child_by_field_name("value")
358 .is_none_or(|value| value.id() != attribute.id())
359 {
360 writes.push(write(if is_binding_target(attribute) {
365 PythonRuntimeWriteTarget::Container
366 } else {
367 PythonRuntimeWriteTarget::ContainerEscape
368 }));
369 return;
370 }
371 if is_binding_target(parent) {
372 writes.push(write(PythonRuntimeWriteTarget::Key));
373 return;
374 }
375 let Some(subscript) = parent.child_by_field_name("subscript") else {
376 writes.push(write(PythonRuntimeWriteTarget::ContainerEscape));
377 return;
378 };
379 let access = access_key(subscript, source);
380 let read_range = node_source_range(parent);
381 let execution = match enclosing_execution_context(attribute) {
382 Some(function) => {
383 PythonRuntimeExecutionContext::DeferredFunction(node_source_range(function))
384 }
385 None => PythonRuntimeExecutionContext::Module,
386 };
387 let accessor = execution_context_accessor_coverage(attribute, module_root, read_range, budget);
388 reads.push(PythonRuntimeRead {
389 binding_name: node_text(path.root, source).to_owned(),
390 root_name: path.root_module.clone(),
391 container: path.container.clone(),
392 access,
393 range: read_range,
394 root_range: node_source_range(path.root),
395 container_range: node_source_range(attribute),
396 candidate_anchor: node_source_range(attribute),
397 key_range: node_source_range(subscript),
398 lexical_resolution: PythonRuntimeRootResolution::ImportedModule,
399 mutation: PythonRuntimeMutationEvidence::UnknownEffects,
401 execution,
402 accessor,
403 });
404}
405
406fn access_key(subscript: Node<'_>, source: &str) -> PythonRuntimeAccessKey {
408 if let Some(literal) = python_plain_string_literal(subscript, source) {
409 return PythonRuntimeAccessKey::Property(literal.to_owned());
410 }
411 match subscript.kind() {
412 "integer" => node_text(subscript, source)
413 .parse::<u128>()
414 .map_or(PythonRuntimeAccessKey::Unsupported, |index| {
415 PythonRuntimeAccessKey::Index(index)
416 }),
417 "slice" | "tuple" | "list_splat" => PythonRuntimeAccessKey::Unsupported,
418 "string" | "concatenated_string" => PythonRuntimeAccessKey::Unsupported,
419 _ => PythonRuntimeAccessKey::Dynamic,
420 }
421}
422
423fn is_binding_target(node: Node<'_>) -> bool {
426 let mut current = node;
427 while let Some(parent) = current.parent() {
428 match parent.kind() {
429 "assignment" | "augmented_assignment" | "for_statement" | "for_in_clause" => {
430 return parent
431 .child_by_field_name("left")
432 .is_some_and(|left| left.id() == current.id());
433 }
434 "delete_statement" => return true,
435 "pattern_list" | "tuple_pattern" | "list_pattern" | "list_splat_pattern" => {
436 current = parent;
437 }
438 _ => return false,
439 }
440 }
441 false
442}
443
444fn enclosing_execution_context<'tree>(node: Node<'tree>) -> Option<Node<'tree>> {
447 let mut current = node;
448 while let Some(parent) = current.parent() {
449 if matches!(parent.kind(), "function_definition" | "lambda")
450 && parent
451 .child_by_field_name("body")
452 .is_some_and(|body| covers(body, current))
453 {
454 return Some(parent);
455 }
456 current = parent;
457 }
458 None
459}
460
461fn execution_context_accessor_coverage(
468 attribute: Node<'_>,
469 module_root: Node<'_>,
470 read_range: Range,
471 budget: &mut PythonRuntimeBudget,
472) -> PythonRuntimeAccessorCoverage {
473 let scope = enclosing_execution_context(attribute)
474 .and_then(|function| function.child_by_field_name("body"))
475 .unwrap_or(module_root);
476 let mut stack = vec![scope];
477 while let Some(node) = stack.pop() {
478 if !budget.visit() {
479 return PythonRuntimeAccessorCoverage::UnknownAccessorEffects;
480 }
481 if matches!(node.kind(), "function_definition" | "lambda") {
482 let mut cursor = node.walk();
485 stack.extend(node.named_children(&mut cursor).filter(|child| {
486 Some(child.id()) != node.child_by_field_name("body").map(|body| body.id())
487 }));
488 continue;
489 }
490 if node.kind() == "call" && can_precede(node_source_range(node), read_range, node) {
491 return PythonRuntimeAccessorCoverage::UnknownAccessorEffects;
492 }
493 let mut cursor = node.walk();
494 stack.extend(node.named_children(&mut cursor));
495 }
496 PythonRuntimeAccessorCoverage::NoKnownAccessorEffects
497}
498
499fn can_precede(hazard: Range, read_range: Range, hazard_node: Node<'_>) -> bool {
505 if hazard.start_byte < read_range.start_byte {
506 return true;
507 }
508 let mut current = hazard_node;
509 while let Some(parent) = current.parent() {
510 if matches!(parent.kind(), "for_statement" | "while_statement")
511 && parent.start_byte() <= read_range.start_byte
512 && read_range.end_byte <= parent.end_byte()
513 {
514 return true;
515 }
516 current = parent;
517 }
518 false
519}
520
521fn escaping_module_identifier(
527 node: Node<'_>,
528 bindings: &[ModuleBinding],
529 module_root: Node<'_>,
530 source: &str,
531 budget: &mut PythonRuntimeBudget,
532) -> Option<String> {
533 let name = node_text(node, source);
534 let module = binding_root_module(node, name, bindings)?;
535 if root_resolution(node, &module, bindings, module_root, source, budget)
539 != PythonRuntimeRootResolution::ImportedModule
540 {
541 return None;
542 }
543 let parent = node.parent()?;
544 if parent.kind() == "attribute"
545 && parent
546 .child_by_field_name("object")
547 .is_some_and(|object| object.id() == node.id())
548 {
549 return None;
550 }
551 if matches!(
552 parent.kind(),
553 "import_statement" | "import_from_statement" | "aliased_import" | "dotted_name"
554 ) {
555 return None;
556 }
557 Some(module)
558}
559
560fn binding_root_module(
563 reference: Node<'_>,
564 name: &str,
565 bindings: &[ModuleBinding],
566) -> Option<String> {
567 let offset = reference.start_byte();
568 let mut module = None;
569 for binding in bindings {
570 if binding.local_name != name
571 || offset < binding.scope_start_byte
572 || offset > binding.scope_end_byte
573 {
574 continue;
575 }
576 match &module {
577 None => module = Some(binding.root_module.clone()),
578 Some(existing) if *existing == binding.root_module => {}
579 Some(_) => return None,
580 }
581 }
582 module.filter(|module| {
583 PROCESS_INPUT_CONTAINERS
584 .iter()
585 .any(|(reviewed, _)| reviewed == module)
586 })
587}
588
589fn root_resolution(
596 root: Node<'_>,
597 expected_module: &str,
598 bindings: &[ModuleBinding],
599 module_root: Node<'_>,
600 source: &str,
601 budget: &mut PythonRuntimeBudget,
602) -> PythonRuntimeRootResolution {
603 let name = node_text(root, source);
604 match python_name_resolution_at(name, root, source) {
605 PythonLexicalNameResolution::Local | PythonLexicalNameResolution::Nonlocal => {
606 return if bindings
610 .iter()
611 .any(|binding| binding.local_name == name && binding.function_scoped)
612 {
613 PythonRuntimeRootResolution::Unknown
614 } else {
615 PythonRuntimeRootResolution::LexicallyBound
616 };
617 }
618 PythonLexicalNameResolution::Global | PythonLexicalNameResolution::Unbound => {}
619 }
620 match module_scope_binds_name_through_imports(module_root, name, source, budget) {
621 Some(true) => {}
622 Some(false) => return PythonRuntimeRootResolution::LexicallyBound,
623 None => return PythonRuntimeRootResolution::Unknown,
624 }
625 match binding_root_module(root, name, bindings) {
626 Some(module) if module == expected_module => PythonRuntimeRootResolution::ImportedModule,
627 Some(_) => PythonRuntimeRootResolution::LexicallyBound,
628 None => PythonRuntimeRootResolution::Unknown,
629 }
630}
631
632fn module_scope_binds_name_through_imports(
638 module_root: Node<'_>,
639 name: &str,
640 source: &str,
641 budget: &mut PythonRuntimeBudget,
642) -> Option<bool> {
643 let mut imported = false;
644 let mut stack = vec![module_root];
645 while let Some(node) = stack.pop() {
646 if !budget.visit() {
647 return None;
648 }
649 if node.kind() == "wildcard_import" {
650 return None;
651 }
652 for binding in python_direct_scope_bindings_bounded(node, source, || budget.visit())? {
653 if node_text(binding.declaration, source) != name {
654 continue;
655 }
656 match binding.kind {
657 PythonDirectScopeBindingKind::Import => imported = true,
658 _ => return Some(false),
659 }
660 }
661 let excluded_body = matches!(
664 node.kind(),
665 "class_definition" | "function_definition" | "lambda"
666 )
667 .then(|| node.child_by_field_name("body").map(|body| body.id()))
668 .flatten();
669 let mut cursor = node.walk();
670 stack.extend(
671 node.named_children(&mut cursor)
672 .filter(|child| Some(child.id()) != excluded_body),
673 );
674 }
675 Some(imported)
676}
677
678fn module_bindings(root: Node<'_>, source: &str) -> Vec<ModuleBinding> {
685 let mut bindings = Vec::new();
686 for binding in python_import_bindings_from_tree(root, source) {
687 let Some(root_module) = binding.qualified_name.split('.').next() else {
688 continue;
689 };
690 if !PROCESS_INPUT_CONTAINERS
691 .iter()
692 .any(|(reviewed, _)| *reviewed == root_module)
693 {
694 continue;
695 }
696 if binding.qualified_name != root_module && binding.consumed_attributes == 0 {
698 continue;
699 }
700 bindings.push(ModuleBinding {
701 local_name: binding.local_name.clone(),
702 root_module: root_module.to_owned(),
703 scope_start_byte: binding.scope_start_byte,
704 scope_end_byte: binding.scope_end_byte,
705 function_scoped: binding.is_function_scoped(),
706 });
707 }
708 bindings
709}
710
711fn covers(container: Node<'_>, node: Node<'_>) -> bool {
712 container.start_byte() <= node.start_byte() && node.end_byte() <= container.end_byte()
713}
714
715fn node_source_range(node: Node<'_>) -> Range {
716 Range {
717 start_byte: node.start_byte(),
718 end_byte: node.end_byte(),
719 start_line: node.start_position().row,
720 end_line: node.end_position().row,
721 }
722}
723
724fn node_text<'a>(node: Node<'_>, source: &'a str) -> &'a str {
725 node.utf8_text(source.as_bytes())
726 .expect("a tree-sitter node range is valid UTF-8 source")
727}
728
729#[cfg(test)]
730mod tests {
731 use super::*;
732
733 fn facts(source: &str) -> PythonRuntimeReadFacts {
734 let mut parser = tree_sitter::Parser::new();
735 parser
736 .set_language(&tree_sitter_python::LANGUAGE.into())
737 .expect("python grammar loads");
738 let tree = parser.parse(source, None).expect("fixture parses");
739 extract_python_runtime_reads(tree.root_node(), source, 64)
740 }
741
742 #[test]
744 fn imported_module_environment_subscript_is_a_pristine_keyed_read() {
745 let source = "import os\n\n\ndef run():\n command = os.environ[\"DFB\"]\n os.system(command)\n";
746 let facts = facts(source);
747 assert!(facts.complete, "{facts:#?}");
748 let [read] = facts.reads.as_slice() else {
749 panic!("exactly one keyed read: {facts:#?}");
750 };
751 assert_eq!(read.root_name, "os");
752 assert_eq!(read.container, "environ");
753 assert_eq!(
754 read.access,
755 PythonRuntimeAccessKey::Property("DFB".to_owned())
756 );
757 assert_eq!(
758 read.lexical_resolution,
759 PythonRuntimeRootResolution::ImportedModule
760 );
761 assert_eq!(read.mutation, PythonRuntimeMutationEvidence::NoKnownWrite);
762 assert_eq!(
763 read.accessor,
764 PythonRuntimeAccessorCoverage::NoKnownAccessorEffects,
765 "the only call runs after the read"
766 );
767 assert!(facts.writes.is_empty(), "{facts:#?}");
768 assert_eq!(
771 &source[read.candidate_anchor.start_byte..read.candidate_anchor.end_byte],
772 "os.environ"
773 );
774 assert_eq!(
775 &source[read.range.start_byte..read.range.end_byte],
776 "os.environ[\"DFB\"]"
777 );
778 }
779
780 #[test]
781 fn imported_module_argv_index_is_a_keyed_read() {
782 let facts = facts("import sys\n\n\ndef run():\n return sys.argv[1]\n");
783 let [read] = facts.reads.as_slice() else {
784 panic!("exactly one keyed read: {facts:#?}");
785 };
786 assert_eq!(read.root_name, "sys");
787 assert_eq!(read.container, "argv");
788 assert_eq!(read.access, PythonRuntimeAccessKey::Index(1));
789 }
790
791 #[test]
794 fn a_shadowed_module_name_is_not_a_read() {
795 let facts = facts("import os\n\n\ndef run(os):\n return os.environ[\"DFB\"]\n");
796 assert!(facts.reads.is_empty(), "{facts:#?}");
797 }
798
799 #[test]
800 fn a_local_container_name_is_not_a_read() {
801 let facts =
802 facts("def run():\n environ = {\"DFB\": \"x\"}\n return environ[\"DFB\"]\n");
803 assert!(facts.reads.is_empty(), "{facts:#?}");
804 assert!(facts.writes.is_empty(), "{facts:#?}");
805 }
806
807 #[test]
808 fn an_unrelated_module_spelled_os_is_not_a_read() {
809 let facts = facts("from mypkg import os\n\n\ndef run():\n return os.environ[\"DFB\"]\n");
810 assert!(facts.reads.is_empty(), "{facts:#?}");
811 }
812
813 #[test]
814 fn a_dynamic_key_stays_a_read_with_no_static_key() {
815 let facts = facts("import os\n\n\ndef run(name):\n return os.environ[name]\n");
816 let [read] = facts.reads.as_slice() else {
817 panic!("one read: {facts:#?}");
818 };
819 assert_eq!(read.access, PythonRuntimeAccessKey::Dynamic);
820 }
821
822 #[test]
823 fn a_container_write_refutes_every_read_of_that_container() {
824 let facts = facts(
825 "import os\n\n\ndef run():\n os.environ[\"A\"] = \"1\"\n return os.environ[\"B\"]\n",
826 );
827 let [read] = facts.reads.as_slice() else {
828 panic!("one read: {facts:#?}");
829 };
830 assert_eq!(read.mutation, PythonRuntimeMutationEvidence::KnownWrite);
831 let [write] = facts.writes.as_slice() else {
832 panic!("one write: {facts:#?}");
833 };
834 assert_eq!(write.target, PythonRuntimeWriteTarget::Key);
835 assert_eq!(write.container.as_deref(), Some("environ"));
836 }
837
838 #[test]
839 fn a_container_method_receiver_refutes_the_pristine_claim() {
840 let facts = facts(
841 "import os\n\n\ndef run():\n os.environ.update({})\n return os.environ[\"B\"]\n",
842 );
843 let [read] = facts.reads.as_slice() else {
844 panic!("one read: {facts:#?}");
845 };
846 assert_eq!(
847 read.mutation,
848 PythonRuntimeMutationEvidence::KnownWriteAndUnknownEffects
849 );
850 }
851
852 #[test]
853 fn a_module_object_escape_refutes_every_container() {
854 let facts =
855 facts("import os\n\n\ndef run(sink):\n sink(os)\n return os.environ[\"B\"]\n");
856 let [read] = facts.reads.as_slice() else {
857 panic!("one read: {facts:#?}");
858 };
859 assert_eq!(
860 read.mutation,
861 PythonRuntimeMutationEvidence::KnownWriteAndUnknownEffects
862 );
863 assert!(
864 facts
865 .writes
866 .iter()
867 .any(|write| write.target == PythonRuntimeWriteTarget::ModuleEscape),
868 "{facts:#?}"
869 );
870 }
871
872 #[test]
873 fn a_call_before_the_read_leaves_accessor_coverage_open() {
874 let facts = facts(
875 "import os\n\n\ndef run(prepare):\n prepare()\n return os.environ[\"B\"]\n",
876 );
877 let [read] = facts.reads.as_slice() else {
878 panic!("one read: {facts:#?}");
879 };
880 assert_eq!(
881 read.accessor,
882 PythonRuntimeAccessorCoverage::UnknownAccessorEffects
883 );
884 }
885
886 #[test]
889 fn a_call_in_another_function_does_not_poison_the_read() {
890 let source = "import os\n\n\ndef other(prepare):\n prepare()\n\n\ndef run():\n return os.environ[\"B\"]\n";
891 let facts = facts(source);
892 let [read] = facts.reads.as_slice() else {
893 panic!("one read: {facts:#?}");
894 };
895 assert_eq!(
896 read.accessor,
897 PythonRuntimeAccessorCoverage::NoKnownAccessorEffects
898 );
899 }
900
901 #[test]
902 fn a_call_later_in_the_same_loop_can_precede_the_read() {
903 let source = "import os\n\n\ndef run(items, prepare):\n for item in items:\n value = os.environ[\"B\"]\n prepare(value)\n";
904 let facts = facts(source);
905 let [read] = facts.reads.as_slice() else {
906 panic!("one read: {facts:#?}");
907 };
908 assert_eq!(
909 read.accessor,
910 PythonRuntimeAccessorCoverage::UnknownAccessorEffects
911 );
912 }
913
914 #[test]
915 fn an_aliased_import_keeps_the_absolute_module_identity() {
916 let facts =
917 facts("import os as operating\n\n\ndef run():\n return operating.environ[\"B\"]\n");
918 let [read] = facts.reads.as_slice() else {
919 panic!("one read: {facts:#?}");
920 };
921 assert_eq!(read.binding_name, "operating");
922 assert_eq!(read.root_name, "os");
923 }
924
925 #[test]
928 fn a_submodule_import_agrees_with_its_package_binding() {
929 let facts =
930 facts("import os\nimport os.path\n\n\ndef run():\n return os.environ[\"B\"]\n");
931 let [read] = facts.reads.as_slice() else {
932 panic!("one read: {facts:#?}");
933 };
934 assert_eq!(read.root_name, "os");
935 }
936
937 #[test]
938 fn a_module_scope_rebinding_of_the_import_name_is_not_a_read() {
939 let facts = facts("import os\n\nos = None\n\n\ndef run():\n return os.environ[\"B\"]\n");
940 assert!(facts.reads.is_empty(), "{facts:#?}");
941 }
942}
943
944#[derive(Clone, Debug, Eq, PartialEq)]
946pub enum PythonModuleMemberMutation {
947 Write(Range),
949 Escape(Range),
952 Budget,
954}
955
956pub fn python_module_member_mutation(
979 root: Node<'_>,
980 source: &str,
981 module: &str,
982 member: &str,
983 max_nodes: usize,
984) -> Option<PythonModuleMemberMutation> {
985 let mut budget = PythonRuntimeBudget {
986 remaining: max_nodes,
987 exhausted: false,
988 visited_nodes: 0,
989 };
990 let locals: Vec<String> = python_import_bindings_from_tree(root, source)
991 .into_iter()
992 .filter(|binding| {
993 binding.qualified_name == module
994 || binding
995 .qualified_name
996 .split('.')
997 .next()
998 .is_some_and(|root_module| {
999 root_module == module && binding.consumed_attributes > 0
1000 })
1001 })
1002 .map(|binding| binding.local_name)
1003 .collect();
1004 let mut stack = vec![root];
1007 while let Some(node) = stack.pop() {
1008 if !budget.visit() {
1009 return Some(PythonModuleMemberMutation::Budget);
1010 }
1011 if node.kind() == "identifier" {
1012 let name = node_text(node, source);
1013 if name == "sys"
1016 && node.parent().is_some_and(|parent| {
1017 parent.kind() == "attribute"
1018 && parent
1019 .child_by_field_name("attribute")
1020 .is_some_and(|attribute| node_text(attribute, source) == "modules")
1021 })
1022 {
1023 return Some(PythonModuleMemberMutation::Escape(node_source_range(node)));
1024 }
1025 if locals.iter().any(|local| local == name)
1026 && !matches!(
1027 python_name_resolution_at(name, node, source),
1028 PythonLexicalNameResolution::Local | PythonLexicalNameResolution::Nonlocal
1029 )
1030 && let Some(mutation) = module_binding_use(node, source, member)
1031 {
1032 return Some(mutation);
1033 }
1034 }
1035 let mut cursor = node.walk();
1036 stack.extend(node.named_children(&mut cursor));
1037 }
1038 None
1039}
1040
1041fn module_binding_use(
1043 node: Node<'_>,
1044 source: &str,
1045 member: &str,
1046) -> Option<PythonModuleMemberMutation> {
1047 let Some(parent) = node.parent() else {
1048 return Some(PythonModuleMemberMutation::Escape(node_source_range(node)));
1049 };
1050 if matches!(
1051 parent.kind(),
1052 "import_statement" | "import_from_statement" | "aliased_import" | "dotted_name"
1053 ) {
1054 return None;
1055 }
1056 if parent.kind() != "attribute"
1057 || parent
1058 .child_by_field_name("object")
1059 .is_none_or(|object| object.id() != node.id())
1060 {
1061 return Some(PythonModuleMemberMutation::Escape(node_source_range(node)));
1062 }
1063 let Some(attribute) = parent.child_by_field_name("attribute") else {
1064 return Some(PythonModuleMemberMutation::Escape(node_source_range(
1065 parent,
1066 )));
1067 };
1068 if node_text(attribute, source) != member {
1069 return None;
1070 }
1071 is_binding_target(parent).then(|| PythonModuleMemberMutation::Write(node_source_range(parent)))
1072}
1073
1074#[cfg(test)]
1075mod module_member_tests {
1076 use super::*;
1077
1078 fn mutation(source: &str, member: &str) -> Option<PythonModuleMemberMutation> {
1079 let mut parser = tree_sitter::Parser::new();
1080 parser
1081 .set_language(&tree_sitter_python::LANGUAGE.into())
1082 .expect("python grammar loads");
1083 let tree = parser.parse(source, None).expect("fixture parses");
1084 python_module_member_mutation(tree.root_node(), source, "os", member, 4096)
1085 }
1086
1087 #[test]
1088 fn an_ordinary_call_replaces_nothing() {
1089 assert_eq!(
1090 mutation("import os\n\n\ndef run(c):\n os.system(c)\n", "system"),
1091 None
1092 );
1093 }
1094
1095 #[test]
1096 fn a_member_assignment_refuses_that_member() {
1097 let found = mutation(
1098 "import os\n\n\ndef patch(fake):\n os.system = fake\n",
1099 "system",
1100 );
1101 assert!(
1102 matches!(found, Some(PythonModuleMemberMutation::Write(_))),
1103 "{found:?}"
1104 );
1105 assert_eq!(
1107 mutation(
1108 "import os\n\n\ndef patch(fake):\n os.getcwd = fake\n",
1109 "system"
1110 ),
1111 None
1112 );
1113 }
1114
1115 #[test]
1116 fn a_module_object_escape_refuses_every_member() {
1117 for source in [
1118 "import os\n\n\ndef leak(sink):\n sink(os)\n",
1119 "import os\n\n\nalias = os\n",
1120 "import os\nimport functools\n\n\ndef patch():\n setattr(os, \"system\", None)\n",
1121 ] {
1122 let found = mutation(source, "system");
1123 assert!(
1124 matches!(found, Some(PythonModuleMemberMutation::Escape(_))),
1125 "{source}: {found:?}"
1126 );
1127 }
1128 }
1129
1130 #[test]
1131 fn a_module_table_read_refuses_every_member_even_without_the_import() {
1132 let found = mutation(
1133 "import sys\n\n\ndef patch(fake):\n sys.modules[\"os\"].system = fake\n",
1134 "system",
1135 );
1136 assert!(
1137 matches!(found, Some(PythonModuleMemberMutation::Escape(_))),
1138 "{found:?}"
1139 );
1140 }
1141
1142 #[test]
1143 fn a_file_that_never_binds_the_module_proves_the_absence() {
1144 assert_eq!(mutation("def run():\n return 1\n", "system"), None);
1145 }
1146}
1147
1148#[cfg(test)]
1149mod shadowing_tests {
1150 use super::*;
1151
1152 fn facts(source: &str) -> PythonRuntimeReadFacts {
1153 let mut parser = tree_sitter::Parser::new();
1154 parser
1155 .set_language(&tree_sitter_python::LANGUAGE.into())
1156 .expect("python grammar loads");
1157 let tree = parser.parse(source, None).expect("fixture parses");
1158 extract_python_runtime_reads(tree.root_node(), source, 64)
1159 }
1160
1161 #[test]
1164 fn a_parameter_spelling_the_module_name_is_not_an_escape() {
1165 let source = "import os\n\n\ndef read_shadowed(os):\n return os.environ[\"A\"]\n\n\ndef read_real():\n return os.environ[\"B\"]\n";
1166 let facts = facts(source);
1167 assert!(facts.writes.is_empty(), "{facts:#?}");
1168 let [read] = facts.reads.as_slice() else {
1169 panic!("only the unshadowed read is a candidate: {facts:#?}");
1170 };
1171 assert_eq!(read.mutation, PythonRuntimeMutationEvidence::NoKnownWrite);
1172 assert_eq!(
1173 &source[read.key_range.start_byte..read.key_range.end_byte],
1174 "\"B\""
1175 );
1176 }
1177
1178 #[test]
1179 fn a_parameter_spelling_the_module_name_does_not_refuse_a_member() {
1180 let source =
1181 "import os\n\n\ndef helper(os):\n return os\n\n\ndef run(c):\n os.system(c)\n";
1182 let mut parser = tree_sitter::Parser::new();
1183 parser
1184 .set_language(&tree_sitter_python::LANGUAGE.into())
1185 .expect("python grammar loads");
1186 let tree = parser.parse(source, None).expect("fixture parses");
1187 assert_eq!(
1188 python_module_member_mutation(tree.root_node(), source, "os", "system", 4096),
1189 None
1190 );
1191 }
1192}