Support for related traces

Author: Hasenpfote

malloc_tracer/tracer.py

@@ -5,6 +5,8 @@
 import math
 import contextlib
 import textwrap
+import os
+import linecache
 from tracemalloc import start, take_snapshot, stop, Filter
 
 
@@ -142,6 +144,40 @@ def extract_dependencies(obj):
     return collector.dependencies
 
 
+class TraceRecord(object):
+
+    def __init__(self, filepath, lineno, size):
+        self._filepath = filepath
+        self._lineno = lineno
+        self._size = size
+
+    @property
+    def filepath(self):
+        return self._filepath
+
+    @property
+    def short_filepath(self):
+        return os.sep.join(self._filepath.split(os.sep)[-2:])
+
+    @property
+    def lineno(self):
+        return self._lineno
+
+    @property
+    def size(self):
+        return self._size
+
+    @property
+    def human_readable_size(self):
+        return bytes_to_hrf(self._size)
+
+    @property
+    def line(self):
+        line = linecache.getline(self._filepath, self._lineno).rstrip()
+        #linecache.clearcache()
+        return line
+
+
 class Tracer(object):
     '''Tracing malloc that occurs inside a function or method.
 
@@ -235,7 +271,8 @@ def _take_snapshot(
     def trace(
         self,
         target_args=None,
-        setup='pass'
+        setup='pass',
+        enable_related_traces=False
     ):
         '''Display the trace result.
 
@@ -248,28 +285,63 @@ def trace(
             target_args=target_args,
             setup=setup
         )
-        snapshot = snapshot.filter_traces([Filter(True, DUMMY_SRC_NAME), ])
-        stats = snapshot.statistics('lineno')
 
-        total = 0
-        detected_lines = dict()
+        traces_records = dict()
+        stats = snapshot.statistics('lineno')
         for stat in stats:
             frame = stat.traceback[0]
-            detected_lines[str(frame.lineno)] = stat.size
-            total += stat.size
+            key = frame.filename
+            traces_record = traces_records.get(key)
+            if traces_record is None:
+                traces_records[key] = dict()
+                traces_record = traces_records.get(key)
+
+            traces_record[frame.lineno] = TraceRecord(
+                filepath=frame.filename,
+                lineno=frame.lineno,
+                size=stat.size
+            )
 
+        # for DUMMY_SRC_NAME.
+        print('<< Target traces >>')
         print('File "{}"'.format(self._filepath))
-        print('Total {}(raw {} B)'.format(bytes_to_hrf(total), total))
         print('Line #    Trace         Line Contents')
         print('=' * (24+80))
 
+        traces_record = traces_records.get(DUMMY_SRC_NAME)
         source_text = ''.join(self._source_lines).rstrip()
-
         for lineno, line in enumerate(source_text.split(sep='\n'), 1):
-            size = detected_lines.get(str(lineno))
-            trace = ' ' * 10 if size is None else bytes_to_hrf(size)
+            if traces_record is None:
+                trace = ' ' * 10
+            else:
+                tr = traces_record.get(lineno)
+                trace = ' ' * 10 if tr is None else tr.human_readable_size
+
             print('{lineno:6d}    {trace:10s}    {contents}'.format(
                 lineno=self._lineno + lineno - 1,
                 trace=trace,
                 contents=line
             ))
+
+        # for others.
+        if enable_related_traces:
+            related_traces_records = {k: v for k, v in traces_records.items() if k != DUMMY_SRC_NAME}
+            if len(related_traces_records) > 0:
+                print()
+                print('<< Related traces >>')
+                print('Line #    Trace         Line Contents')
+                print('=' * (24+80))
+
+                for filepath, traces_record in sorted(related_traces_records.items()):
+                    print('File "{}"'.format(filepath))
+                    for lineno, tr in sorted(traces_record.items()):
+                        print('{lineno:6d}    {trace:10s}    {contents}'.format(
+                            lineno=tr.lineno,
+                            trace=tr.human_readable_size,
+                            contents=tr.line
+                        ))
+
+        # Total allocated size.
+        total = sum(stat.size for stat in stats)
+        print()
+        print('Total allocated size: {} (raw {} B)'.format(bytes_to_hrf(total), total))

test/test_tracer.py

@@ -69,6 +69,11 @@ def test_function(self):
                 target_args=dict(base=2, num=100),
                 setup='import math as mathematics'
             )
+            tracer.trace(
+                target_args=dict(base=2, num=100),
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_function_with_default_args(self):
         tracer = Tracer(
@@ -80,6 +85,10 @@ def test_function_with_default_args(self):
             tracer.trace(
                 setup='import math as mathematics'
             )
+            tracer.trace(
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_function_with_auto_resolve_dependencies(self):
         tracer = Tracer(
@@ -88,6 +97,7 @@ def test_function_with_auto_resolve_dependencies(self):
         )
         with contextlib.redirect_stdout(None):
             tracer.trace()
+            tracer.trace(enable_related_traces=True)
 
     def test_method(self):
         instance = Klass(2, 100)
@@ -99,6 +109,10 @@ def test_method(self):
             tracer.trace(
                 setup='import math as mathematics'
             )
+            tracer.trace(
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_method_with_default_args(self):
         instance = Klass2()
@@ -111,6 +125,10 @@ def test_method_with_default_args(self):
             tracer.trace(
                 setup='import math as mathematics'
             )
+            tracer.trace(
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_method_with_auto_resolve_dependencies(self):
         instance = Klass2()
@@ -120,6 +138,7 @@ def test_method_with_auto_resolve_dependencies(self):
         )
         with contextlib.redirect_stdout(None):
             tracer.trace()
+            tracer.trace(enable_related_traces=True)
 
     def test_static_method(self):
         tracer = Tracer(
@@ -131,6 +150,11 @@ def test_static_method(self):
                 target_args=dict(base=2, num=100),
                 setup='import math as mathematics'
             )
+            tracer.trace(
+                target_args=dict(base=2, num=100),
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_static_method_with_default_args(self):
         tracer = Tracer(
@@ -143,6 +167,11 @@ def test_static_method_with_default_args(self):
                 target_args=dict(base=2, num=100),
                 setup='import math as mathematics'
             )
+            tracer.trace(
+                target_args=dict(base=2, num=100),
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_static_method_with_auto_resolve_dependencies(self):
         tracer = Tracer(
@@ -153,6 +182,10 @@ def test_static_method_with_auto_resolve_dependencies(self):
             tracer.trace(
                 target_args=dict(base=2, num=100),
             )
+            tracer.trace(
+                target_args=dict(base=2, num=100),
+                enable_related_traces=True
+            )
 
     def test_class_method(self):
         tracer = Tracer(
@@ -164,6 +197,11 @@ def test_class_method(self):
                 target_args=dict(base=1),
                 setup='import math as mathematics',
             )
+            tracer.trace(
+                target_args=dict(base=1),
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_class_method_with_default_args(self):
         tracer = Tracer(
@@ -176,6 +214,11 @@ def test_class_method_with_default_args(self):
                 target_args=dict(base=1),
                 setup='import math as mathematics',
             )
+            tracer.trace(
+                target_args=dict(base=1),
+                setup='import math as mathematics',
+                enable_related_traces=True
+            )
 
     def test_class_method_with_auto_resolve_dependencies(self):
         tracer = Tracer(
@@ -186,3 +229,7 @@ def test_class_method_with_auto_resolve_dependencies(self):
             tracer.trace(
                 target_args=dict(base=1),
             )
+            tracer.trace(
+                target_args=dict(base=1),
+                enable_related_traces=True
+            )