| 637 | |
| 638 | |
| 639 | class CodeMap(dict): |
| 640 | def __init__(self, include_children, backend): |
| 641 | self.include_children = include_children |
| 642 | self._toplevel = [] |
| 643 | self.backend = backend |
| 644 | |
| 645 | def add(self, code, toplevel_code=None): |
| 646 | if code in self: |
| 647 | return |
| 648 | |
| 649 | if toplevel_code is None: |
| 650 | filename = code.co_filename |
| 651 | if filename.endswith((".pyc", ".pyo")): |
| 652 | filename = filename[:-1] |
| 653 | if not os.path.exists(filename): |
| 654 | print('ERROR: Could not find file ' + filename) |
| 655 | if filename.startswith(("ipython-input", "<ipython-input")): |
| 656 | print( |
| 657 | "NOTE: %mprun can only be used on functions defined in" |
| 658 | " physical files, and not in the IPython environment.") |
| 659 | return |
| 660 | |
| 661 | toplevel_code = code |
| 662 | (sub_lines, start_line) = inspect.getsourcelines(code) |
| 663 | linenos = range(start_line, |
| 664 | start_line + len(sub_lines)) |
| 665 | self._toplevel.append((filename, code, linenos)) |
| 666 | self[code] = {} |
| 667 | else: |
| 668 | self[code] = self[toplevel_code] |
| 669 | |
| 670 | for subcode in filter(inspect.iscode, code.co_consts): |
| 671 | self.add(subcode, toplevel_code=toplevel_code) |
| 672 | |
| 673 | def trace(self, code, lineno, prev_lineno): |
| 674 | memory = _get_memory(-1, self.backend, include_children=self.include_children, |
| 675 | filename=code.co_filename) |
| 676 | prev_value = self[code].get(lineno, None) |
| 677 | previous_memory = prev_value[1] if prev_value else 0 |
| 678 | previous_inc = prev_value[0] if prev_value else 0 |
| 679 | |
| 680 | prev_line_value = self[code].get(prev_lineno, None) if prev_lineno else None |
| 681 | prev_line_memory = prev_line_value[1] if prev_line_value else 0 |
| 682 | occ_count = self[code][lineno][2] + 1 if lineno in self[code] else 1 |
| 683 | self[code][lineno] = ( |
| 684 | previous_inc + (memory - prev_line_memory), |
| 685 | max(memory, previous_memory), |
| 686 | occ_count, |
| 687 | ) |
| 688 | |
| 689 | def items(self): |
| 690 | """Iterate on the toplevel code blocks.""" |
| 691 | for (filename, code, linenos) in self._toplevel: |
| 692 | measures = self[code] |
| 693 | if not measures: |
| 694 | continue # skip if no measurement |
| 695 | line_iterator = ((line, measures.get(line)) for line in linenos) |
| 696 | yield (filename, line_iterator) |