| 544 | |
| 545 | @dataclasses.dataclass |
| 546 | class Tracer: |
| 547 | # the modules we should trace calls to |
| 548 | calls_to_modules: List[str] |
| 549 | # the modules we should trace calls from |
| 550 | calls_from_modules: List[str] |
| 551 | |
| 552 | def __enter__(self): |
| 553 | sys.settrace(self) |
| 554 | |
| 555 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 556 | sys.settrace(None) |
| 557 | |
| 558 | def should_trace(self, *values) -> bool: |
| 559 | for value in values: |
| 560 | if isinstance(value, types.BuiltinMethodType): |
| 561 | # if this was a method defined in C, use the instance as the value |
| 562 | value = value.__self__ |
| 563 | module = getmodulename(value) or getmodulename(type(value)) |
| 564 | |
| 565 | if not module: |
| 566 | warnings.warn(f"Cannot get module of {value}") |
| 567 | continue |
| 568 | if any(module.startswith(mod) for mod in self.calls_to_modules): |
| 569 | return True |
| 570 | return False |
| 571 | |
| 572 | def __call__(self, frame, event, arg) -> Optional[Tracer]: |
| 573 | if event == "call": |
| 574 | frame.f_trace_opcodes = True |
| 575 | return self |
| 576 | elif event != "opcode": |
| 577 | return None |
| 578 | |
| 579 | if self.should_trace_frame(frame): |
| 580 | Stack(self, frame)() |
| 581 | return None |
| 582 | |
| 583 | def should_trace_frame(self, frame) -> bool: |
| 584 | # Ignore frames which are not from the `calls_from_module` |
| 585 | try: |
| 586 | frame_module_name = frame.f_globals["__name__"] |
| 587 | except KeyError: |
| 588 | return False |
| 589 | return frame_module_name == "__main__" or any( |
| 590 | frame_module_name == mod or frame_module_name.startswith(mod + ".") |
| 591 | for mod in self.calls_from_modules |
| 592 | ) |
| 593 | |
| 594 | |
| 595 | def setup(): |
no outgoing calls