Add ignore_non_files

gaogaotiantian · Sep 17, 2020 · 8a2aa66 · 8a2aa66
1 parent 0ac4f4b
commit 8a2aa66
Show file tree

Hide file tree

Showing 9 changed files with 69 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -47,18 +47,18 @@ python3 my_script.py arg1 arg2
 You can simply use VizTracer as 
 
 ```
-python3 -m viztracer my_script.py arg1 arg2
-# OR
 viztracer my_script arg1 arg2
+# OR
+python3 -m viztracer my_script.py arg1 arg2
 ```
 
 which will generate a ```result.html``` file in the directory you run this command. Open it in browser and there's your result.
 
 You can also generate ```json``` file or ```gz``` file and load it with [chrome://tracing/](chrome://tracing/) or [perfetto](https://ui.perfetto.dev/). ```gz``` file is especially helpful when your trace file is large
 
 ```
-python3 -m viztracer -o result.json my_script.py arg1 arg2
-python3 -m viztracer -o result.json.gz my_script.py arg1 arg2
+viztracer -o result.json my_script.py arg1 arg2
+viztracer -o result.json.gz my_script.py arg1 arg2
 ```
 
 ### Inline

diff --git a/docs/source/advanced_usage.rst b/docs/source/advanced_usage.rst
@@ -56,6 +56,18 @@ OR
     
     tracer = VizTracer(ignore_c_function=True)
 
+You can ask VizTracer not to trace any functions that are not in a valid file(mostly import stuff) using ``ignore_non_file``
+
+.. code-block:: 
+
+    viztracer --ignore_non_file my_script.py
+
+OR
+
+.. code-block:: python
+    
+    tracer = VizTracer(ignore_non_file=True)
+
 It's possible that you want to ignore some arbitrary functions and their descendants. You can do it using ``@ignore_function`` decorator
 
 .. code-block:: python

diff --git a/docs/source/viztracer.rst b/docs/source/viztracer.rst
@@ -8,6 +8,7 @@ VizTracer
                  include_files=None,\
                  exclude_files=None,\
                  ignore_c_function=False,\
+                 ignore_non_file=False,\
                  log_return_value=False,\
                  log_function_args=False,\
                  log_print=False,\
@@ -114,6 +115,18 @@ VizTracer
 
             python -m viztracer --ignore_c_function
 
+    .. py:attribute:: ignore_non_file
+        :type: boolean
+        :value: False
+
+        Whether trace functions from invalid files(mostly import stuff)
+
+        Setting it to ``True`` is equivalent to 
+
+        .. code-block::
+
+            python -m viztracer --ignore_non_file
+
     .. py:attribute:: log_return_value 
         :type: boolean
         :value: False

diff --git a/src/viztracer/main.py b/src/viztracer/main.py
@@ -31,6 +31,8 @@ def main():
                         help="specify the only files(directories) you want to include from tracing. Can't be used with --exclude_files")
     parser.add_argument("--ignore_c_function", action="store_true", default=False,
                         help="ignore all c functions including most builtin functions and libraries")
+    parser.add_argument("--ignore_non_file", action="store_true", default=False,
+                        help="ignore all functions that are not in a vaild file(like import)")
     parser.add_argument("--log_return_value", action="store_true", default=False,
                         help="log return value of the function in the report")
     parser.add_argument("--log_print", action="store_true", default=False,
@@ -141,6 +143,7 @@ def main():
         exclude_files=options.exclude_files,
         include_files=options.include_files,
         ignore_c_function=options.ignore_c_function,
+        ignore_non_file=options.ignore_non_file,
         log_return_value=options.log_return_value,
         log_function_args=options.log_function_args,
         log_print=options.log_print,

diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c
@@ -364,6 +364,16 @@ snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg
             }
         }
 
+        if (CHECK_FLAG(self->check_flags, SNAPTRACE_IGNORE_NON_FILE)) {
+            if (is_python && is_call) {
+                PyObject* file_name = frame->f_code->co_filename;
+                if (startswith(PyUnicode_AsUTF8(file_name), "<")) {
+                    info->ignore_stack_depth += 1;
+                    return 0;
+                }
+            }
+        }
+
         if (is_call) {
             // If it's a call, we need a new node, and we need to update the stack
             if (!info->stack_top->next) {
@@ -725,25 +735,27 @@ static PyObject*
 snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
 {
     static char* kwlist[] = {"verbose", "lib_file_path", "max_stack_depth", 
-            "include_files", "exclude_files", "ignore_c_function", "log_return_value", 
-            "novdb", "log_function_args",
+            "include_files", "exclude_files", "ignore_c_function", "ignore_non_file",
+            "log_return_value", "novdb", "log_function_args",
             NULL};
     int kw_verbose = -1;
     int kw_max_stack_depth = 0;
     char* kw_lib_file_path = NULL;
     PyObject* kw_include_files = NULL;
     PyObject* kw_exclude_files = NULL;
     int kw_ignore_c_function = -1;
+    int kw_ignore_non_file = -1;
     int kw_log_return_value = -1;
     int kw_novdb = -1;
     int kw_log_function_args = -1;
-    if (!PyArg_ParseTupleAndKeywords(args, kw, "|isiOOpppp", kwlist, 
+    if (!PyArg_ParseTupleAndKeywords(args, kw, "|isiOOppppp", kwlist, 
             &kw_verbose,
             &kw_lib_file_path,
             &kw_max_stack_depth,
             &kw_include_files,
             &kw_exclude_files,
             &kw_ignore_c_function,
+            &kw_ignore_non_file,
             &kw_log_return_value,
             &kw_novdb,
             &kw_log_function_args)) {
@@ -776,6 +788,12 @@ snaptrace_config(TracerObject* self, PyObject* args, PyObject* kw)
         UNSET_FLAG(self->check_flags, SNAPTRACE_IGNORE_C_FUNCTION);
     }
 
+    if (kw_ignore_non_file == 1) {
+        SET_FLAG(self->check_flags, SNAPTRACE_IGNORE_NON_FILE);
+    } else if (kw_ignore_non_file == 0) {
+        UNSET_FLAG(self->check_flags, SNAPTRACE_IGNORE_NON_FILE);
+    }
+
     if (kw_log_return_value == 1) {
         SET_FLAG(self->check_flags, SNAPTRACE_LOG_RETURN_VALUE);
     } else if (kw_log_return_value == 0) {

diff --git a/src/viztracer/modules/snaptrace.h b/src/viztracer/modules/snaptrace.h
@@ -11,6 +11,7 @@
 #define SNAPTRACE_LOG_RETURN_VALUE (1 << 4)
 #define SNAPTRACE_NOVDB (1 << 5)
 #define SNAPTRACE_LOG_FUNCTION_ARGS (1 << 6)
+#define SNAPTRACE_IGNORE_NON_FILE (1 << 7)
 
 #define SET_FLAG(reg, flag) ((reg) |= (flag))
 #define UNSET_FLAG(reg, flag) ((reg) &= (~(flag)))

diff --git a/src/viztracer/tracer.py b/src/viztracer/tracer.py
@@ -17,6 +17,7 @@ def __init__(self,
                  include_files=None,
                  exclude_files=None,
                  ignore_c_function=False,
+                 ignore_non_file=False,
                  log_return_value=False,
                  log_function_args=False,
                  log_print=False,
@@ -33,6 +34,7 @@ def __init__(self,
         self.include_files = include_files
         self.exclude_files = exclude_files
         self.ignore_c_function = ignore_c_function
+        self.ignore_non_file = ignore_non_file
         self.log_return_value = log_return_value
         self.log_print = log_print
         self.novdb = novdb
@@ -147,6 +149,7 @@ def start(self):
             include_files=self.include_files,
             exclude_files=self.exclude_files,
             ignore_c_function=self.ignore_c_function,
+            ignore_non_file=self.ignore_non_file,
             log_return_value=self.log_return_value,
             novdb=self.novdb,
             log_function_args=self.log_function_args

diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py
@@ -19,6 +19,7 @@ def __init__(self,
                  include_files=None,
                  exclude_files=None,
                  ignore_c_function=False,
+                 ignore_non_file=False,
                  log_return_value=False,
                  log_function_args=False,
                  log_print=False,
@@ -32,6 +33,7 @@ def __init__(self,
                 include_files=include_files,
                 exclude_files=exclude_files,
                 ignore_c_function=ignore_c_function,
+                ignore_non_file=ignore_non_file,
                 log_return_value=log_return_value,
                 log_print=log_print,
                 novdb=novdb,

diff --git a/tests/test_tracer.py b/tests/test_tracer.py
@@ -157,6 +157,16 @@ def test_ignore_c_function(self):
         entries = tracer.parse()
         self.assertEqual(entries, 0)
 
+    def test_ignore_non_file(self):
+        tracer = _VizTracer(ignore_non_file=True)
+        tracer.start()
+        import random
+        lst = []
+        lst.append(1)
+        tracer.stop()
+        entries = tracer.parse()
+        self.assertEqual(entries, 1)
+
     def test_log_return_value(self):
         tracer = _VizTracer()
         tracer.start()