Add a feature to filter output by include and exclude patterns

Author: Hasenpfote

malloc_tracer/tracer.py

@@ -6,6 +6,7 @@
 import contextlib
 import textwrap
 import linecache
+import fnmatch
 from enum import Enum
 from tracemalloc import start, take_snapshot, stop, Filter
 
@@ -295,7 +296,9 @@ def trace(
         self,
         target_args=None,
         setup='pass',
-        related_traces_output_mode=RelatedTracesOutputMode.NONE
+        related_traces_output_mode=RelatedTracesOutputMode.NONE,
+        include_patterns=None,
+        exclude_patterns=None
     ):
         '''Display the trace result.
 
@@ -304,16 +307,31 @@ def trace(
             setup (str): Run-time dependencies.
                 This parameter is ignored if enable_auto_resolve is enabled.
             related_traces_output_mode (:class:`RelatedTracesOutputMode`):
+            include_patterns (set): Specify patterns of file paths to include in the output.
+            exclude_patterns (set): Specify patterns of file paths to exclude in the output.
         '''
         snapshot = self._take_snapshot(
             target_args=target_args,
             setup=setup
         )
 
+        if include_patterns is None:
+            include_patterns = set('*')
+
+        if exclude_patterns is None:
+            exclude_patterns = set()
+
         recorder = TraceRecorder()
         stats = snapshot.statistics('lineno')
         for stat in stats:
             frame = stat.traceback[0]
+
+            if frame.filename != DUMMY_SRC_NAME:
+                if not any(fnmatch.fnmatch(frame.filename, pattern) for pattern in include_patterns):
+                    continue
+                if any(fnmatch.fnmatch(frame.filename, pattern) for pattern in exclude_patterns):
+                    continue
+
             recorder.add_trace(
                 filepath=frame.filename,
                 lineno=frame.lineno,
@@ -447,7 +465,9 @@ def trace(
     enable_auto_resolve=True,
     ctime_setup='pass',
     rtime_setup='pass',
-    related_traces_output_mode=RelatedTracesOutputMode.NONE
+    related_traces_output_mode=RelatedTracesOutputMode.NONE,
+    include_patterns=None,
+    exclude_patterns=None
 ):
     '''Convenience function to create Tracer object and call trace method.'''
     tracer = Tracer(
@@ -458,5 +478,7 @@ def trace(
     tracer.trace(
         target_args=target_args,
         setup=rtime_setup,
-        related_traces_output_mode=related_traces_output_mode
+        related_traces_output_mode=related_traces_output_mode,
+        include_patterns=include_patterns,
+        exclude_patterns=exclude_patterns
     )

test/test_tracer.py

@@ -280,3 +280,23 @@ def test_class_method_with_auto_resolve_dependencies(self):
                 target_args=dict(base=1),
                 related_traces_output_mode=RelatedTracesOutputMode.IN_DESCENDING_ORDER
             )
+
+    def test_include_and_exclude_patterns(self):
+        tracer = Tracer(
+            function,
+        )
+        with contextlib.redirect_stdout(None):
+            tracer.trace(
+                target_args=dict(base=2, num=100),
+                setup='import math as mathematics',
+                related_traces_output_mode=RelatedTracesOutputMode.FOR_EACH_FILE,
+                include_patterns={'*.py'},
+                exclude_patterns={'*/site-packages/*'}
+            )
+            tracer.trace(
+                target_args=dict(base=2, num=100),
+                setup='import math as mathematics',
+                related_traces_output_mode=RelatedTracesOutputMode.IN_DESCENDING_ORDER,
+                include_patterns={'*.py'},
+                exclude_patterns={'*/site-packages/*'}
+            )