From 61b5e0f39741d00064e4d185df6ccaca6109b068 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Fri, 4 Aug 2023 16:04:42 +0000
Subject: [PATCH 01/18] Fix log_sparse for context manager

---
 src/viztracer/viztracer.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py
index e104da7e..15d798cb 100644
--- a/src/viztracer/viztracer.py
+++ b/src/viztracer/viztracer.py
@@ -124,11 +124,13 @@ def init_kwargs(self) -> Dict:
         }
 
     def __enter__(self) -> "VizTracer":
-        self.start()
+        if not self.log_sparse:
+            self.start()
         return self
 
     def __exit__(self, type, value, trace) -> None:
-        self.stop()
+        if not self.log_sparse:
+            self.stop()
         self.save()
         self.terminate()
 

From 7039e192fec6b3fa08f6641d3a7c8c63d532ace5 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Fri, 4 Aug 2023 16:32:09 +0000
Subject: [PATCH 02/18] Allow to define log_sparse functions outside of
 decorator

---
 src/viztracer/decorator.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index 8067e5a3..7431be9a 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -58,16 +58,13 @@ def wrapper(*args, **kwargs) -> Any:
 
 
 def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0) -> Callable:
-    tracer = get_tracer()
-    if tracer is None or not tracer.log_sparse:
-        if func is None:
-            return lambda f: f
-        return func
-
     if func is None:
         def inner(dec_func: Callable) -> Callable:
             @functools.wraps(dec_func)
             def wrapper(*args, **kwargs) -> Any:
+                tracer = get_tracer()
+                if tracer is None or not tracer.log_sparse:
+                    return dec_func(*args, **kwargs)
                 assert isinstance(tracer, VizTracer)
                 if not tracer.enable:
                     orig_max_stack_depth = tracer.max_stack_depth
@@ -84,6 +81,9 @@ def wrapper(*args, **kwargs) -> Any:
     else:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:
+            tracer = get_tracer()
+            if tracer is None or not tracer.log_sparse:
+                return func(*args, **kwargs)
             assert callable(func)
             assert isinstance(tracer, VizTracer)
             start = tracer.getts()

From 526e70002ad00d46df7963201039bf30c6850721 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Tue, 8 Aug 2023 15:28:17 +0000
Subject: [PATCH 03/18] Add dynamic_tracer_check parameter

---
 src/viztracer/decorator.py | 80 ++++++++++++++++++++++++--------------
 1 file changed, 51 insertions(+), 29 deletions(-)

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index 7431be9a..5cf83ee2 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -56,39 +56,52 @@ def wrapper(*args, **kwargs) -> Any:
         return inner(method)
     return inner
 
-
-def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0) -> Callable:
-    if func is None:
-        def inner(dec_func: Callable) -> Callable:
-            @functools.wraps(dec_func)
-            def wrapper(*args, **kwargs) -> Any:
-                tracer = get_tracer()
-                if tracer is None or not tracer.log_sparse:
-                    return dec_func(*args, **kwargs)
-                assert isinstance(tracer, VizTracer)
-                if not tracer.enable:
-                    orig_max_stack_depth = tracer.max_stack_depth
-                    tracer.max_stack_depth = stack_depth
-                    tracer.start()
-                    ret = dec_func(*args, **kwargs)
-                    tracer.stop()
-                    tracer.max_stack_depth = orig_max_stack_depth
-                    return ret
-                else:
-                    return dec_func(*args, **kwargs)
-            return wrapper
-        return inner
-    else:
+def _log_sparse_wrapper(func: Optional[Callable] = None, stack_depth: int = 0,
+                        dynamic_tracer_check: bool = False) -> Callable:
+    if not dynamic_tracer_check:
+        tracer = get_tracer()
+        if tracer is None or not tracer.log_sparse:
+            if func is None:
+                return lambda f: f
+            return func
+
+    if stack_depth > 0:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:
-            tracer = get_tracer()
-            if tracer is None or not tracer.log_sparse:
+            local_tracer = get_tracer() if dynamic_tracer_check else tracer
+
+            if local_tracer is None or not local_tracer.log_sparse:
+                if func is None:
+                    return lambda f: f
+                return func
+
+            assert isinstance(local_tracer, VizTracer)
+            if not local_tracer.enable:
+                orig_max_stack_depth = local_tracer.max_stack_depth
+                local_tracer.max_stack_depth = stack_depth
+                local_tracer.start()
+                ret = func(*args, **kwargs)
+                local_tracer.stop()
+                local_tracer.max_stack_depth = orig_max_stack_depth
+                return ret
+            else:
                 return func(*args, **kwargs)
+        return wrapper
+    else:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs) -> Any:
+            local_tracer = get_tracer() if dynamic_tracer_check else tracer
+
+            if local_tracer is None or not local_tracer.log_sparse:
+                if func is None:
+                    return lambda f: f
+                return func
+
             assert callable(func)
-            assert isinstance(tracer, VizTracer)
-            start = tracer.getts()
+            assert isinstance(local_tracer, VizTracer)
+            start = local_tracer.getts()
             ret = func(*args, **kwargs)
-            dur = tracer.getts() - start
+            dur = local_tracer.getts() - start
             code = func.__code__
             raw_data = {
                 "ph": "X",
@@ -97,7 +110,16 @@ def wrapper(*args, **kwargs) -> Any:
                 "dur": dur,
                 "cat": "FEE",
             }
-            tracer.add_raw(raw_data)
+            local_tracer.add_raw(raw_data)
             return ret
 
         return wrapper
+
+
+def _log_sparse_wrapper_with_args(stack_depth: int = 0, dynamic_tracer_check: bool = False) -> Callable:
+    return functools.partial(_log_sparse_wrapper, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
+
+def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0, dynamic_tracer_check: bool = False) -> Callable:
+    if func is not None:
+        return _log_sparse_wrapper(func=func, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
+    return _log_sparse_wrapper_with_args(stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)

From 82680a270aacaaa82ff12173b6eb061bf50b2903 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Wed, 9 Aug 2023 09:59:31 +0000
Subject: [PATCH 04/18] Update documentation

---
 docs/source/filter.rst | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/docs/source/filter.rst b/docs/source/filter.rst
index d8a78435..94af74cf 100644
--- a/docs/source/filter.rst
+++ b/docs/source/filter.rst
@@ -79,14 +79,14 @@ Ignore C Function
 
 By default, VizTracer will trace both python and C functions. You can turn off tracing C functions by
 
-.. code-block:: 
+.. code-block::
 
     viztracer --ignore_c_function my_script.py
 
 OR
 
 .. code-block:: python
-    
+
     tracer = VizTracer(ignore_c_function=True)
 
 Since most of the builtin functions(like ``append`` or ``len``) are C functions which are frequently called,
@@ -98,14 +98,14 @@ Ignore Non File
 
 You can ask VizTracer not to trace any functions that are not in a valid file(mostly import stuff) using ``ignore_frozen``
 
-.. code-block:: 
+.. code-block::
 
     viztracer --ignore_frozen my_script.py
 
 OR
 
 .. code-block:: python
-    
+
     tracer = VizTracer(ignore_frozen=True)
 
 
@@ -128,7 +128,7 @@ Log Sparse
 ----------
 
 You can make VizTracer log only certain functions using ``--log_sparse``. This is helpful when you are only interested in the time spent on
-specific functions for a big picture on larger projects. 
+specific functions for a big picture on larger projects.
 
 First, you need to add decorator ``@log_sparse`` on the function you want to log
 
@@ -151,13 +151,17 @@ First, you need to add decorator ``@log_sparse`` on the function you want to log
 
 Then just call viztracer with ``--log_sparse``
 
-.. code-block:: 
+.. code-block::
 
     viztracer --log_sparse your_script.py
-    
+
 When you are using ``--log_sparse``, due to the nature of the recording, some advanced features may not work with it.
 
 You can leave ``@log_sparse`` as it is when you are not running the script with VizTracer. It will be like a no-op
 
 If you want to log a piece of code, rather than a full function, please check :ref:`duration_event_label`. Duration Event
-is compatible with ``log_sparse``
\ No newline at end of file
+is compatible with ``log_sparse``
+
+To use ``@log_sparse`` in conjunction with a context manager, you must define decorating functions within the created
+context, or set the `dynamic_tracer_check=True`` argument of decorator. The second option leads to runtime checks,
+so it increases the overhead.

From 7fb9dcb104efc4ce896fe88810ccf6df6534e74f Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Wed, 9 Aug 2023 10:48:17 +0000
Subject: [PATCH 05/18] Add kwargs to %%viztracer

---
 src/viztracer/cellmagic.py | 40 +++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py
index f92665a5..a7409d09 100644
--- a/src/viztracer/cellmagic.py
+++ b/src/viztracer/cellmagic.py
@@ -1,5 +1,7 @@
 # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
+import re
+import inspect
 
 try:
     from IPython.core.magic import (Magics, cell_magic,  # type: ignore
@@ -15,13 +17,15 @@ def viztracer(self, line, cell, local_ns) -> None:
 
             from .viewer import ServerThread
             from .viztracer import VizTracer
+            tracer_kwargs, viewer_kwargs = self.make_kwargs(
+                line, inspect.signature(VizTracer), inspect.signature(ServerThread)
+            )
             code = self.shell.transform_cell(cell)
-            file_path = "./viztracer_report.json"
-            with VizTracer(verbose=0, output_file=file_path):
+            with VizTracer(**tracer_kwargs):
                 exec(code, local_ns, local_ns)
 
             def view():  # pragma: no cover
-                server = ServerThread(file_path, once=True)
+                server = ServerThread(**viewer_kwargs)
                 server.start()
                 server.ready.wait()
                 import webbrowser
@@ -32,6 +36,36 @@ def view():  # pragma: no cover
 
             display(button)
 
+        def make_kwargs(self, line, tracer_signature, viewer_signature):
+            line = line.split('#')[0]
+            pattern = r'(\w+=\w+)\s'
+            default_path = "./viztracer_report.json"
+
+            tracer_kwargs = {'verbose': 0, 'output_file': default_path}
+            viewer_kwargs = {'path': default_path, 'once': True}
+
+            for item in re.split(pattern, line):
+                if item == '':
+                    continue
+                name, value = item.split('=')
+                if name == 'output_file':
+                    tracer_kwargs['output_file'] = value
+                    viewer_kwargs['path'] = value
+                elif name in viewer_signature.parameters:
+                    annotation = viewer_signature.parameters[name].annotation
+                    kwargs = viewer_kwargs
+                elif name in tracer_signature.parameters:
+                    annotation = tracer_signature.parameters[name].annotation
+                    kwargs = tracer_kwargs
+                else:
+                    raise ValueError(f'{name} is not a parameter of VizTracer or ServerThread')
+
+                if value == 'None':
+                    kwargs[name] = None
+                else:
+                    kwargs[name] = annotation(value)
+            return tracer_kwargs, viewer_kwargs
+
 except ImportError:  # pragma: no cover
     pass
 

From 346c6213d6510d5372e9eb381d08bc9c4583dbe4 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Wed, 9 Aug 2023 12:00:56 +0000
Subject: [PATCH 06/18] Fix log_sparse=False mode

---
 src/viztracer/decorator.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index 5cf83ee2..a02f5942 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -71,9 +71,7 @@ def wrapper(*args, **kwargs) -> Any:
             local_tracer = get_tracer() if dynamic_tracer_check else tracer
 
             if local_tracer is None or not local_tracer.log_sparse:
-                if func is None:
-                    return lambda f: f
-                return func
+                return func(*args, **kwargs)
 
             assert isinstance(local_tracer, VizTracer)
             if not local_tracer.enable:
@@ -93,9 +91,7 @@ def wrapper(*args, **kwargs) -> Any:
             local_tracer = get_tracer() if dynamic_tracer_check else tracer
 
             if local_tracer is None or not local_tracer.log_sparse:
-                if func is None:
-                    return lambda f: f
-                return func
+                return func(*args, **kwargs)
 
             assert callable(func)
             assert isinstance(local_tracer, VizTracer)

From 7e91b53362d1a92e8fbe6e7db40ef058023feaea Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Thu, 10 Aug 2023 09:04:08 +0000
Subject: [PATCH 07/18] Update docstrings

---
 docs/source/basic_usage.rst |  6 +++++
 docs/source/decorator.rst   | 47 +++++++++++++++++++++++++++++++------
 docs/source/filter.rst      |  8 +++++++
 3 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/docs/source/basic_usage.rst b/docs/source/basic_usage.rst
index fda4caf9..c21fbf4d 100644
--- a/docs/source/basic_usage.rst
+++ b/docs/source/basic_usage.rst
@@ -110,6 +110,12 @@ If you are using Jupyter, you can use viztracer cell magics.
     %%viztracer
     # Your code after
 
+.. code-block:: python
+
+    # you can define arguments of VizTracer in magic
+    %%viztracer port=8888
+    # Your code after
+
 A ``Show VizTracer Report`` button will appear after the cell and you can click it to view the results
 
 Display Report
diff --git a/docs/source/decorator.rst b/docs/source/decorator.rst
index a1b9c9f2..34b06088 100644
--- a/docs/source/decorator.rst
+++ b/docs/source/decorator.rst
@@ -3,7 +3,7 @@ Decorator
 
 .. py:decorator:: ignore_function
 
-    ``@ignore_function`` can tell VizTracer to skip on functions you specified. 
+    ``@ignore_function`` can tell VizTracer to skip on functions you specified.
 
     .. code-block:: python
 
@@ -12,21 +12,21 @@ Decorator
         def function_you_want_to_ignore():
             # function body
 
-        # You can specify tracer if no global tracer is registered 
+        # You can specify tracer if no global tracer is registered
         @ignore_function(tracer=tracer)
         def function_you_want_to_ignore():
             # function body
 
 .. py:decorator:: trace_and_save(method=None, output_dir="./", **viztracer_kwargs)
-    
+
     :param function method: trick to make both ``@trace_and_save`` and ``@trace_and_save(**kwargs)`` work
     :param str output_dir: output directory you want to put your logs in
     :param dict viztracer_kwargs: kwargs for VizTracer
 
-    ``@trace_and_save`` can be used to trace a certain function and save the result as the program goes. 
-    This can be very helpful when you are running a very long program and just want to keep recording 
+    ``@trace_and_save`` can be used to trace a certain function and save the result as the program goes.
+    This can be very helpful when you are running a very long program and just want to keep recording
     something periodically. You won't drain your memory and the parsing/dumping will be done in a new process,
-    which can minimize the performance impact to your main process. 
+    which can minimize the performance impact to your main process.
 
     You can pass any argument you want to ``VizTracer`` by giving it to the decorator
 
@@ -39,4 +39,37 @@ Decorator
         # this works as well
         @trace_and_save
         def function_you_want_to_trace():
-            # function body
\ No newline at end of file
+            # function body
+
+.. py:decorator:: log_sparse(func=None, stack_depth=0, dynamic_tracer_check=False)
+
+    You can make VizTracer log only certain functions using ``--log_sparse`` mode.
+
+    :param function func: callable to decorate
+    :param int stack_depth: log the function and its descendants with a limit stack depth
+    :param bool dynamic_tracer_check: run time check of tracer
+
+    .. code-block:: python
+
+        from viztracer import log_sparse
+
+        # @log_sparse will only log this function
+        @log_sparse
+        def function_you_want_to_log():
+            # function body
+
+        # @log_sparse(stack_depth=5) will log this function and its descendants
+        # with a limit stack depth of 5
+        # Nested @log_sparse with stack_depth won't work
+        # (only the outermost function and its stack will be logged)
+        @log_sparse(stack_depth=5)
+        def function_you_want_to_log():
+            # function body
+
+        # Use dynamic_tracer_check=True if you use tracer as a context manager (or with %%viztracer).
+        @log_sparse(dynamic_tracer_check=True)
+        def function_you_want_to_log():
+            # function body
+
+        with VizTracer(log_sparse=True):
+            function_you_want_to_log()
diff --git a/docs/source/filter.rst b/docs/source/filter.rst
index 94af74cf..18632792 100644
--- a/docs/source/filter.rst
+++ b/docs/source/filter.rst
@@ -149,6 +149,14 @@ First, you need to add decorator ``@log_sparse`` on the function you want to log
     def function_you_want_to_log():
         # function body
 
+    # Use dynamic_tracer_check=True if you use tracer as a context manager (or with %%viztracer).
+    @log_sparse(dynamic_tracer_check=True)
+    def function_you_want_to_log():
+        # function body
+
+    with VizTracer(log_sparse=True):
+        function_you_want_to_log()
+
 Then just call viztracer with ``--log_sparse``
 
 .. code-block::

From f9d60b8d355e559e061c04ca11a7c8d44b701d00 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Thu, 10 Aug 2023 09:47:05 +0000
Subject: [PATCH 08/18] Add strings processing for magic

---
 src/viztracer/cellmagic.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py
index a7409d09..52edb2f7 100644
--- a/src/viztracer/cellmagic.py
+++ b/src/viztracer/cellmagic.py
@@ -49,21 +49,25 @@ def make_kwargs(self, line, tracer_signature, viewer_signature):
                     continue
                 name, value = item.split('=')
                 if name == 'output_file':
+                    value = value.strip(' ').strip('\'"')
                     tracer_kwargs['output_file'] = value
                     viewer_kwargs['path'] = value
-                elif name in viewer_signature.parameters:
-                    annotation = viewer_signature.parameters[name].annotation
-                    kwargs = viewer_kwargs
-                elif name in tracer_signature.parameters:
-                    annotation = tracer_signature.parameters[name].annotation
-                    kwargs = tracer_kwargs
                 else:
-                    raise ValueError(f'{name} is not a parameter of VizTracer or ServerThread')
+                    if name in viewer_signature.parameters:
+                        annotation = viewer_signature.parameters[name].annotation
+                        kwargs = viewer_kwargs
+                    elif name in tracer_signature.parameters:
+                        annotation = tracer_signature.parameters[name].annotation
+                        kwargs = tracer_kwargs
+                    else:
+                        raise ValueError(f'{name} is not a parameter of VizTracer or ServerThread')
 
-                if value == 'None':
-                    kwargs[name] = None
-                else:
-                    kwargs[name] = annotation(value)
+                    if value == 'None':
+                        kwargs[name] = None
+                    elif annotation is str:
+                        kwargs[name] = value.strip(' ').strip('\'"')
+                    else:
+                        kwargs[name] = annotation(value)
             return tracer_kwargs, viewer_kwargs
 
 except ImportError:  # pragma: no cover

From d602260e259019342098e9da52874336292e9824 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Mon, 14 Aug 2023 16:50:25 +0000
Subject: [PATCH 09/18] Changes in magic args

---
 docs/source/basic_usage.rst |  2 +-
 src/viztracer/cellmagic.py  | 82 +++++++++++++++++++------------------
 2 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/docs/source/basic_usage.rst b/docs/source/basic_usage.rst
index c21fbf4d..1af25514 100644
--- a/docs/source/basic_usage.rst
+++ b/docs/source/basic_usage.rst
@@ -113,7 +113,7 @@ If you are using Jupyter, you can use viztracer cell magics.
 .. code-block:: python
 
     # you can define arguments of VizTracer in magic
-    %%viztracer port=8888
+    %%viztracer -p 8888
     # Your code after
 
 A ``Show VizTracer Report`` button will appear after the cell and you can click it to view the results
diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py
index 52edb2f7..85ec1bf7 100644
--- a/src/viztracer/cellmagic.py
+++ b/src/viztracer/cellmagic.py
@@ -1,14 +1,38 @@
 # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
-import re
-import inspect
 
 try:
     from IPython.core.magic import (Magics, cell_magic,  # type: ignore
                                     magics_class, needs_local_scope)
+    from IPython.core.magic_arguments import (argument, magic_arguments, parse_argstring)
 
     @magics_class
     class VizTracerMagics(Magics):
+        @magic_arguments()
+        @argument("--port", "-p", default=9001, type=int,
+                  help="specify the port vizviewer will use")
+        @argument("--output_file", default="./viztracer_report.json",
+                  help="output file path. End with .json or .html or .gz")
+        @argument("--max_stack_depth", type=int, default=-1,
+                  help="maximum stack depth you want to trace.")
+        @argument("--verbose", "-v", default=0, type=int,
+                  help="verbosity level")
+        @argument("--ignore_c_function", action="store_true", default=False,
+                  help="ignore all c functions including most builtin functions and libraries")
+        @argument("--ignore_frozen", action="store_true", default=False,
+                  help="ignore all functions that are frozen(like import)")
+        @argument("--log_exit", action="store_true", default=False,
+                  help="log functions in exit functions like atexit")
+        @argument("--log_func_args", action="store_true", default=False,
+                  help="log all function arguments, this will introduce large overhead")
+        @argument("--log_print", action="store_true", default=False,
+                  help="replace all print() function to adding an event to the result")
+        @argument("--log_sparse", action="store_true", default=False,
+                  help="log only selected functions with @log_sparse")
+        @argument("--log_func_entry", nargs="*", default=None,
+                  help="log entry of the function with specified names")
+        @argument("--log_exception", action="store_true", default=False,
+                  help="log all exception when it's raised")
         @needs_local_scope
         @cell_magic
         def viztracer(self, line, cell, local_ns) -> None:
@@ -17,15 +41,27 @@ def viztracer(self, line, cell, local_ns) -> None:
 
             from .viewer import ServerThread
             from .viztracer import VizTracer
-            tracer_kwargs, viewer_kwargs = self.make_kwargs(
-                line, inspect.signature(VizTracer), inspect.signature(ServerThread)
-            )
+            options = parse_argstring(self.viztracer, line)
             code = self.shell.transform_cell(cell)
+            file_path = options.output_file
+            tracer_kwargs = {
+                "output_file": file_path,
+                "verbose": options.verbose,
+                "max_stack_depth": options.max_stack_depth,
+                "ignore_c_function": options.ignore_c_function,
+                "log_exit": options.log_exit,
+                "ignore_frozen": options.ignore_frozen,
+                "log_func_args": options.log_func_args,
+                "log_print": options.log_print,
+                "log_sparse": options.log_sparse,
+                "log_func_entry": options.log_func_entry,
+                "log_exception": options.log_exception
+            }
             with VizTracer(**tracer_kwargs):
                 exec(code, local_ns, local_ns)
 
             def view():  # pragma: no cover
-                server = ServerThread(**viewer_kwargs)
+                server = ServerThread(file_path, port=options.port, once=True)
                 server.start()
                 server.ready.wait()
                 import webbrowser
@@ -36,40 +72,6 @@ def view():  # pragma: no cover
 
             display(button)
 
-        def make_kwargs(self, line, tracer_signature, viewer_signature):
-            line = line.split('#')[0]
-            pattern = r'(\w+=\w+)\s'
-            default_path = "./viztracer_report.json"
-
-            tracer_kwargs = {'verbose': 0, 'output_file': default_path}
-            viewer_kwargs = {'path': default_path, 'once': True}
-
-            for item in re.split(pattern, line):
-                if item == '':
-                    continue
-                name, value = item.split('=')
-                if name == 'output_file':
-                    value = value.strip(' ').strip('\'"')
-                    tracer_kwargs['output_file'] = value
-                    viewer_kwargs['path'] = value
-                else:
-                    if name in viewer_signature.parameters:
-                        annotation = viewer_signature.parameters[name].annotation
-                        kwargs = viewer_kwargs
-                    elif name in tracer_signature.parameters:
-                        annotation = tracer_signature.parameters[name].annotation
-                        kwargs = tracer_kwargs
-                    else:
-                        raise ValueError(f'{name} is not a parameter of VizTracer or ServerThread')
-
-                    if value == 'None':
-                        kwargs[name] = None
-                    elif annotation is str:
-                        kwargs[name] = value.strip(' ').strip('\'"')
-                    else:
-                        kwargs[name] = annotation(value)
-            return tracer_kwargs, viewer_kwargs
-
 except ImportError:  # pragma: no cover
     pass
 

From fc5468798cffb3b3697c8d618d8a168f4c0f0d96 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Wed, 16 Aug 2023 09:04:20 +0000
Subject: [PATCH 10/18] Update tests; Add setignorestackcounter

---
 src/viztracer/decorator.py        | 72 +++++++++++++++++----------
 src/viztracer/modules/snaptrace.c | 22 +++++++-
 src/viztracer/tracer.py           |  3 ++
 tests/test_logsparse.py           | 83 +++++++++++++++++++++++++++++++
 4 files changed, 152 insertions(+), 28 deletions(-)

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index a02f5942..6fd246e3 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -56,13 +56,12 @@ def wrapper(*args, **kwargs) -> Any:
         return inner(method)
     return inner
 
-def _log_sparse_wrapper(func: Optional[Callable] = None, stack_depth: int = 0,
+
+def _log_sparse_wrapper(func: Optional[Callable], stack_depth: int = 0,
                         dynamic_tracer_check: bool = False) -> Callable:
     if not dynamic_tracer_check:
         tracer = get_tracer()
         if tracer is None or not tracer.log_sparse:
-            if func is None:
-                return lambda f: f
             return func
 
     if stack_depth > 0:
@@ -70,11 +69,12 @@ def _log_sparse_wrapper(func: Optional[Callable] = None, stack_depth: int = 0,
         def wrapper(*args, **kwargs) -> Any:
             local_tracer = get_tracer() if dynamic_tracer_check else tracer
 
-            if local_tracer is None or not local_tracer.log_sparse:
+            if local_tracer is None:
                 return func(*args, **kwargs)
-
             assert isinstance(local_tracer, VizTracer)
-            if not local_tracer.enable:
+
+            if (not dynamic_tracer_check and not local_tracer.enable) or \
+                (dynamic_tracer_check and not local_tracer.enable and local_tracer.log_sparse):
                 orig_max_stack_depth = local_tracer.max_stack_depth
                 local_tracer.max_stack_depth = stack_depth
                 local_tracer.start()
@@ -82,40 +82,58 @@ def wrapper(*args, **kwargs) -> Any:
                 local_tracer.stop()
                 local_tracer.max_stack_depth = orig_max_stack_depth
                 return ret
+            elif dynamic_tracer_check and local_tracer.enable and not local_tracer.log_sparse:
+                # The call is made from the module inside, so if `trace_self=False` it will be ignored.
+                # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
+                # recover it
+                prev_ignore_stack = local_tracer.setignorestackcounter(0)
+                ret = func(*args, **kwargs)
+                local_tracer.setignorestackcounter(prev_ignore_stack)
+                return ret
             else:
                 return func(*args, **kwargs)
+
         return wrapper
     else:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:
             local_tracer = get_tracer() if dynamic_tracer_check else tracer
 
-            if local_tracer is None or not local_tracer.log_sparse:
+            if local_tracer is None:
                 return func(*args, **kwargs)
 
-            assert callable(func)
             assert isinstance(local_tracer, VizTracer)
-            start = local_tracer.getts()
-            ret = func(*args, **kwargs)
-            dur = local_tracer.getts() - start
-            code = func.__code__
-            raw_data = {
-                "ph": "X",
-                "name": f"{code.co_name} ({code.co_filename}:{code.co_firstlineno})",
-                "ts": start,
-                "dur": dur,
-                "cat": "FEE",
-            }
-            local_tracer.add_raw(raw_data)
-            return ret
 
-        return wrapper
+            if (not dynamic_tracer_check and not local_tracer.enable) or \
+                (dynamic_tracer_check and not local_tracer.enable and local_tracer.log_sparse):
+                start = local_tracer.getts()
+                ret = func(*args, **kwargs)
+                dur = local_tracer.getts() - start
+                code = func.__code__
+                raw_data = {
+                    "ph": "X",
+                    "name": f"{code.co_name} ({code.co_filename}:{code.co_firstlineno})",
+                    "ts": start,
+                    "dur": dur,
+                    "cat": "FEE",
+                }
+                local_tracer.add_raw(raw_data)
+                return ret
+            elif dynamic_tracer_check and local_tracer.enable and not local_tracer.log_sparse:
+                # The call is made from the module inside, so if `trace_self=False` it will be ignored.
+                # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
+                # recover it
+                prev_ignore_stack = local_tracer.setignorestackcounter(0)
+                ret = func(*args, **kwargs)
+                local_tracer.setignorestackcounter(prev_ignore_stack)
+                return ret
+            else:
+                return func(*args, **kwargs)
 
+        return wrapper
 
-def _log_sparse_wrapper_with_args(stack_depth: int = 0, dynamic_tracer_check: bool = False) -> Callable:
-    return functools.partial(_log_sparse_wrapper, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
 
 def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0, dynamic_tracer_check: bool = False) -> Callable:
-    if func is not None:
-        return _log_sparse_wrapper(func=func, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
-    return _log_sparse_wrapper_with_args(stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
+    if func is None:
+        return functools.partial(_log_sparse_wrapper, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
+    return _log_sparse_wrapper(func=func, stack_depth=stack_depth, dynamic_tracer_check=dynamic_tracer_check)
diff --git a/src/viztracer/modules/snaptrace.c b/src/viztracer/modules/snaptrace.c
index 8e6449fd..ff662f68 100644
--- a/src/viztracer/modules/snaptrace.c
+++ b/src/viztracer/modules/snaptrace.c
@@ -43,6 +43,7 @@ static PyObject* snaptrace_addfunctionarg(TracerObject* self, PyObject* args);
 static PyObject* snaptrace_getfunctionarg(TracerObject* self, PyObject* args);
 static PyObject* snaptrace_getts(TracerObject* self, PyObject* args);
 static PyObject* snaptrace_setcurrstack(TracerObject* self, PyObject* args);
+static PyObject* snaptrace_setignorestackcounter(TracerObject* self, PyObject* args);
 static void snaptrace_threaddestructor(void* key);
 static struct ThreadInfo* snaptrace_createthreadinfo(TracerObject* self);
 static void log_func_args(struct FunctionNode* node, PyFrameObject* frame);
@@ -218,6 +219,7 @@ static PyMethodDef Tracer_methods[] = {
     {"setcurrstack", (PyCFunction)snaptrace_setcurrstack, METH_VARARGS, "set current stack depth"},
     {"pause", (PyCFunction)snaptrace_pause, METH_VARARGS, "pause profiling"},
     {"resume", (PyCFunction)snaptrace_resume, METH_VARARGS, "resume profiling"},
+    {"setignorestackcounter", (PyCFunction)snaptrace_setignorestackcounter, METH_VARARGS, "reset ignore stack depth"},
     {NULL, NULL, 0, NULL}
 };
 
@@ -301,7 +303,7 @@ snaptrace_tracefunc(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg
                 goto cleanup;
             }
         }
-        
+
         // IMPORTANT: the C function will always be called from our python methods, 
         // so this check is redundant. However, the user should never be allowed 
         // to call our C methods directly! Otherwise C functions will be recorded.
@@ -1414,6 +1416,24 @@ snaptrace_addraw(TracerObject* self, PyObject* args)
     Py_RETURN_NONE;
 }
 
+static PyObject*
+snaptrace_setignorestackcounter(TracerObject* self, PyObject* args)
+{
+    int value = 0;
+    int current_value = 0;
+    struct ThreadInfo* info = get_thread_info(self);
+
+    if (!PyArg_ParseTuple(args, "i", &value)) {
+        printf("Error when parsing arguments!\n");
+        exit(1);
+    }
+
+    current_value = info->ignore_stack_depth;
+    info->ignore_stack_depth = value;
+
+    return Py_BuildValue("i", current_value);
+}
+
 static PyObject*
 snaptrace_getfunctionarg(TracerObject* self, PyObject* args)
 {
diff --git a/src/viztracer/tracer.py b/src/viztracer/tracer.py
index b5a3d9d1..0a00fe44 100644
--- a/src/viztracer/tracer.py
+++ b/src/viztracer/tracer.py
@@ -292,6 +292,9 @@ def enable_thread_tracing(self) -> None:
     def getts(self) -> float:
         return self._tracer.getts()
 
+    def setignorestackcounter(self, value) -> int:
+        return self._tracer.setignorestackcounter(value)
+
     def add_instant(self, name: str, args: Any = None, scope: str = "g") -> None:
         if self.enable:
             if scope not in ["g", "p", "t"]:
diff --git a/tests/test_logsparse.py b/tests/test_logsparse.py
index 534a5329..0a9fa722 100644
--- a/tests/test_logsparse.py
+++ b/tests/test_logsparse.py
@@ -75,6 +75,79 @@ def f(x):
         time.sleep(0.1)
 """
 
+file_context_manager = """
+from viztracer import VizTracer, log_sparse
+
+@log_sparse(dynamic_tracer_check=True)
+def f():
+    return 1
+
+def g():
+    return f()
+
+@log_sparse
+def h():
+    return 2
+
+@log_sparse(dynamic_tracer_check=True, stack_depth=1)
+def q():
+    return 3
+
+if __name__ == "__main__":
+    with VizTracer(output_file="result.json"):
+        assert g() == 1
+        assert h() == 2
+        assert q() == 3
+"""
+
+file_context_manager_logsparse = """
+from viztracer import VizTracer, log_sparse
+
+@log_sparse(dynamic_tracer_check=True)
+def f():
+    return 1
+
+def g():
+    return f()
+
+@log_sparse
+def h():
+    return 2
+
+@log_sparse(dynamic_tracer_check=True, stack_depth=1)
+def q():
+    return 3
+
+if __name__ == "__main__":
+    with VizTracer(output_file="result.json", log_sparse=True):
+        assert g() == 1
+        assert h() == 2
+        assert q() == 3
+"""
+
+file_context_manager_logsparse_stack = """
+from viztracer import VizTracer, log_sparse
+
+@log_sparse(dynamic_tracer_check=True)
+def f():
+    return 1
+
+@log_sparse(dynamic_tracer_check=True, stack_depth=1)
+def g():
+    return f()
+
+@log_sparse(dynamic_tracer_check=True)
+def h():
+    return 2
+
+if __name__ == "__main__":
+    assert g() == 1
+    assert h() == 2
+
+    with VizTracer(output_file="result.json", log_sparse=True):
+        assert g() == 1
+        assert h() == 2
+"""
 
 class TestLogSparse(CmdlineTmpl):
     def test_basic(self):
@@ -118,3 +191,13 @@ def test_multiprocess(self):
                 # coveragepy has some issue with multiprocess pool
                 if not os.getenv("COVERAGE_RUN"):
                     raise e
+
+    def test_context_manager(self):
+        self.template(["python", "cmdline_test.py"], script=file_context_manager,
+                      expected_output_file="result.json", expected_entries=4)
+
+        self.template(["python", "cmdline_test.py"], script=file_context_manager_logsparse,
+                      expected_output_file="result.json", expected_entries=2)
+
+        self.template(["python", "cmdline_test.py"], script=file_context_manager_logsparse_stack,
+                      expected_output_file="result.json", expected_entries=2)

From d1838facc3c97526cb26a462c419ad2f4b68f98d Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Wed, 16 Aug 2023 11:11:49 +0000
Subject: [PATCH 11/18] Remove some arguments from magic

---
 src/viztracer/cellmagic.py | 11 +----------
 1 file changed, 1 insertion(+), 10 deletions(-)

diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py
index 85ec1bf7..71ed0f50 100644
--- a/src/viztracer/cellmagic.py
+++ b/src/viztracer/cellmagic.py
@@ -21,18 +21,12 @@ class VizTracerMagics(Magics):
                   help="ignore all c functions including most builtin functions and libraries")
         @argument("--ignore_frozen", action="store_true", default=False,
                   help="ignore all functions that are frozen(like import)")
-        @argument("--log_exit", action="store_true", default=False,
-                  help="log functions in exit functions like atexit")
         @argument("--log_func_args", action="store_true", default=False,
                   help="log all function arguments, this will introduce large overhead")
         @argument("--log_print", action="store_true", default=False,
                   help="replace all print() function to adding an event to the result")
         @argument("--log_sparse", action="store_true", default=False,
                   help="log only selected functions with @log_sparse")
-        @argument("--log_func_entry", nargs="*", default=None,
-                  help="log entry of the function with specified names")
-        @argument("--log_exception", action="store_true", default=False,
-                  help="log all exception when it's raised")
         @needs_local_scope
         @cell_magic
         def viztracer(self, line, cell, local_ns) -> None:
@@ -49,13 +43,10 @@ def viztracer(self, line, cell, local_ns) -> None:
                 "verbose": options.verbose,
                 "max_stack_depth": options.max_stack_depth,
                 "ignore_c_function": options.ignore_c_function,
-                "log_exit": options.log_exit,
                 "ignore_frozen": options.ignore_frozen,
                 "log_func_args": options.log_func_args,
                 "log_print": options.log_print,
-                "log_sparse": options.log_sparse,
-                "log_func_entry": options.log_func_entry,
-                "log_exception": options.log_exception
+                "log_sparse": options.log_sparse
             }
             with VizTracer(**tracer_kwargs):
                 exec(code, local_ns, local_ns)

From 7e9494508f372a616eb007531e043f029e0f7c7c Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Mon, 23 Oct 2023 16:04:34 +0000
Subject: [PATCH 12/18] Add VizShield

---
 src/viztracer/decorator.py | 51 +++++++++-----------------------------
 src/viztracer/viztracer.py |  4 +++
 2 files changed, 16 insertions(+), 39 deletions(-)

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index 6fd246e3..b0b29fc7 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -7,6 +7,7 @@
 from typing import Any, Callable, Optional
 
 from .viztracer import VizTracer, get_tracer
+from .vizshield import VizShield
 
 
 def ignore_function(method: Optional[Callable] = None, tracer: Optional[VizTracer] = None) -> Callable:
@@ -64,17 +65,16 @@ def _log_sparse_wrapper(func: Optional[Callable], stack_depth: int = 0,
         if tracer is None or not tracer.log_sparse:
             return func
 
-    if stack_depth > 0:
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs) -> Any:
-            local_tracer = get_tracer() if dynamic_tracer_check else tracer
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs) -> Any:
+        local_tracer = get_tracer() if dynamic_tracer_check else tracer
 
-            if local_tracer is None:
-                return func(*args, **kwargs)
-            assert isinstance(local_tracer, VizTracer)
+        if local_tracer is None:
+            return func(*args, **kwargs)
+        assert isinstance(local_tracer, VizTracer)
 
-            if (not dynamic_tracer_check and not local_tracer.enable) or \
-                (dynamic_tracer_check and not local_tracer.enable and local_tracer.log_sparse):
+        if local_tracer.log_sparse and not local_tracer.enable:
+            if stack_depth > 0:
                 orig_max_stack_depth = local_tracer.max_stack_depth
                 local_tracer.max_stack_depth = stack_depth
                 local_tracer.start()
@@ -82,30 +82,7 @@ def wrapper(*args, **kwargs) -> Any:
                 local_tracer.stop()
                 local_tracer.max_stack_depth = orig_max_stack_depth
                 return ret
-            elif dynamic_tracer_check and local_tracer.enable and not local_tracer.log_sparse:
-                # The call is made from the module inside, so if `trace_self=False` it will be ignored.
-                # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
-                # recover it
-                prev_ignore_stack = local_tracer.setignorestackcounter(0)
-                ret = func(*args, **kwargs)
-                local_tracer.setignorestackcounter(prev_ignore_stack)
-                return ret
             else:
-                return func(*args, **kwargs)
-
-        return wrapper
-    else:
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs) -> Any:
-            local_tracer = get_tracer() if dynamic_tracer_check else tracer
-
-            if local_tracer is None:
-                return func(*args, **kwargs)
-
-            assert isinstance(local_tracer, VizTracer)
-
-            if (not dynamic_tracer_check and not local_tracer.enable) or \
-                (dynamic_tracer_check and not local_tracer.enable and local_tracer.log_sparse):
                 start = local_tracer.getts()
                 ret = func(*args, **kwargs)
                 dur = local_tracer.getts() - start
@@ -119,18 +96,14 @@ def wrapper(*args, **kwargs) -> Any:
                 }
                 local_tracer.add_raw(raw_data)
                 return ret
-            elif dynamic_tracer_check and local_tracer.enable and not local_tracer.log_sparse:
+        else:
+            with local_tracer.shield_ignore():
                 # The call is made from the module inside, so if `trace_self=False` it will be ignored.
                 # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
                 # recover it
-                prev_ignore_stack = local_tracer.setignorestackcounter(0)
-                ret = func(*args, **kwargs)
-                local_tracer.setignorestackcounter(prev_ignore_stack)
-                return ret
-            else:
                 return func(*args, **kwargs)
 
-        return wrapper
+    return wrapper
 
 
 def log_sparse(func: Optional[Callable] = None, stack_depth: int = 0, dynamic_tracer_check: bool = False) -> Callable:
diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py
index 15d798cb..ca39e7a1 100644
--- a/src/viztracer/viztracer.py
+++ b/src/viztracer/viztracer.py
@@ -14,6 +14,7 @@
 from .tracer import _VizTracer
 from .vizevent import VizEvent
 from .vizplugin import VizPluginBase, VizPluginManager
+from .vizshield import VizShield
 
 
 # This is the interface of the package. Almost all user should use this
@@ -167,6 +168,9 @@ def log_event(self, event_name: str) -> VizEvent:
         call_frame = sys._getframe(1)
         return VizEvent(self, event_name, call_frame.f_code.co_filename, call_frame.f_lineno)
 
+    def shield_ignore(self):
+        return VizShield(self)
+
     def set_afterfork(self, callback: Callable, *args, **kwargs) -> None:
         self._afterfork_cb = callback
         self._afterfork_args = args

From 3f51ae853761d62d259f79f97bfa2f1ba81bf7c7 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Mon, 23 Oct 2023 16:30:13 +0000
Subject: [PATCH 13/18] Remove verbose parameter of magic

---
 src/viztracer/cellmagic.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py
index cbdba593..ab54751f 100644
--- a/src/viztracer/cellmagic.py
+++ b/src/viztracer/cellmagic.py
@@ -15,8 +15,6 @@ class VizTracerMagics(Magics):
                   help="output file path. End with .json or .html or .gz")
         @argument("--max_stack_depth", type=int, default=-1,
                   help="maximum stack depth you want to trace.")
-        @argument("--verbose", "-v", default=0, type=int,
-                  help="verbosity level")
         @argument("--ignore_c_function", action="store_true", default=False,
                   help="ignore all c functions including most builtin functions and libraries")
         @argument("--ignore_frozen", action="store_true", default=False,
@@ -41,7 +39,7 @@ def viztracer(self, line, cell, local_ns) -> None:
             file_path = options.output_file
             tracer_kwargs = {
                 "output_file": file_path,
-                "verbose": options.verbose,
+                "verbose": 0,
                 "max_stack_depth": options.max_stack_depth,
                 "ignore_c_function": options.ignore_c_function,
                 "ignore_frozen": options.ignore_frozen,

From 0d12d897d86bf5ef4359a51f4da224dc8f2f006b Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Tue, 24 Oct 2023 09:10:59 +0000
Subject: [PATCH 14/18] Add vizshield.py; Update docs

---
 docs/source/basic_usage.rst | 13 ++++++++++++-
 src/viztracer/vizshield.py  | 22 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 src/viztracer/vizshield.py

diff --git a/docs/source/basic_usage.rst b/docs/source/basic_usage.rst
index ff830296..d199afb3 100644
--- a/docs/source/basic_usage.rst
+++ b/docs/source/basic_usage.rst
@@ -116,7 +116,18 @@ If you are using Jupyter, you can use viztracer cell magics.
     %%viztracer -p 8888
     # Your code after
 
-A ``Show VizTracer Report`` button will appear after the cell and you can click it to view the results
+A ``Show VizTracer Report`` button will appear after the cell and you can click it to view the results.
+
+Cell magic ``%%viztracer`` supports some of the command line arguments:
+
+* ``--port``
+* ``--output_file``
+* ``--max_stack_depth``
+* ``--ignore_c_function``
+* ``--ignore_frozen``
+* ``--log_func_args``
+* ``--log_print``
+* ``--log_sparse``
 
 Display Report
 --------------
diff --git a/src/viztracer/vizshield.py b/src/viztracer/vizshield.py
new file mode 100644
index 00000000..e62a5957
--- /dev/null
+++ b/src/viztracer/vizshield.py
@@ -0,0 +1,22 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
+
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .viztracer import VizTracer  # pragma: no cover
+
+
+class VizShield:
+    def __init__(self, tracer: "VizTracer") -> None:
+        self._tracer = tracer
+        self._prev_ignore_stack = None
+
+    def __enter__(self) -> None:
+        if self._tracer.enable and not self._tracer.log_sparse:
+            self._prev_ignore_stack = self._tracer.setignorestackcounter(0)
+
+    def __exit__(self, type, value, trace) -> None:
+        if self._prev_ignore_stack is not None:
+            self._tracer.setignorestackcounter(self._prev_ignore_stack)

From cb1b6f4bef12cb8195cc182659c8d86d3d9950e7 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Mon, 30 Oct 2023 08:35:18 +0000
Subject: [PATCH 15/18] Changes in shield

---
 src/viztracer/decorator.py | 11 ++++++-----
 src/viztracer/vizshield.py | 22 ----------------------
 src/viztracer/viztracer.py |  9 +++++++--
 tests/test_logsparse.py    |  1 +
 4 files changed, 14 insertions(+), 29 deletions(-)
 delete mode 100644 src/viztracer/vizshield.py

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index b0b29fc7..b868fb29 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -7,7 +7,6 @@
 from typing import Any, Callable, Optional
 
 from .viztracer import VizTracer, get_tracer
-from .vizshield import VizShield
 
 
 def ignore_function(method: Optional[Callable] = None, tracer: Optional[VizTracer] = None) -> Callable:
@@ -96,12 +95,14 @@ def wrapper(*args, **kwargs) -> Any:
                 }
                 local_tracer.add_raw(raw_data)
                 return ret
-        else:
+        elif local_tracer.enable and not local_tracer.log_sparse:
+            # The call is made from the module inside, so if `trace_self=False` it will be ignored.
+            # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
+            # recover it
             with local_tracer.shield_ignore():
-                # The call is made from the module inside, so if `trace_self=False` it will be ignored.
-                # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
-                # recover it
                 return func(*args, **kwargs)
+        else:
+            return func(*args, **kwargs)
 
     return wrapper
 
diff --git a/src/viztracer/vizshield.py b/src/viztracer/vizshield.py
deleted file mode 100644
index e62a5957..00000000
--- a/src/viztracer/vizshield.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
-# For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
-
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from .viztracer import VizTracer  # pragma: no cover
-
-
-class VizShield:
-    def __init__(self, tracer: "VizTracer") -> None:
-        self._tracer = tracer
-        self._prev_ignore_stack = None
-
-    def __enter__(self) -> None:
-        if self._tracer.enable and not self._tracer.log_sparse:
-            self._prev_ignore_stack = self._tracer.setignorestackcounter(0)
-
-    def __exit__(self, type, value, trace) -> None:
-        if self._prev_ignore_stack is not None:
-            self._tracer.setignorestackcounter(self._prev_ignore_stack)
diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py
index 8996c719..c9b0a6b9 100644
--- a/src/viztracer/viztracer.py
+++ b/src/viztracer/viztracer.py
@@ -2,6 +2,7 @@
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
 
 import builtins
+from contextlib import contextmanager
 import multiprocessing
 import os
 import signal
@@ -14,7 +15,6 @@
 from .tracer import _VizTracer
 from .vizevent import VizEvent
 from .vizplugin import VizPluginBase, VizPluginManager
-from .vizshield import VizShield
 
 
 # This is the interface of the package. Almost all user should use this
@@ -165,8 +165,13 @@ def log_event(self, event_name: str) -> VizEvent:
         call_frame = sys._getframe(1)
         return VizEvent(self, event_name, call_frame.f_code.co_filename, call_frame.f_lineno)
 
+    @contextmanager
     def shield_ignore(self):
-        return VizShield(self)
+        prev_ignore_stack = self.setignorestackcounter(0)
+        try:
+            yield
+        finally:
+            self.setignorestackcounter(prev_ignore_stack)
 
     def set_afterfork(self, callback: Callable, *args, **kwargs) -> None:
         self._afterfork_cb = callback
diff --git a/tests/test_logsparse.py b/tests/test_logsparse.py
index 0a9fa722..66c2a6ab 100644
--- a/tests/test_logsparse.py
+++ b/tests/test_logsparse.py
@@ -149,6 +149,7 @@ def h():
         assert h() == 2
 """
 
+
 class TestLogSparse(CmdlineTmpl):
     def test_basic(self):
         def check_func(data):

From ca35f7cb9979b1cfb13309d3c6181c958b1af9d7 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Tue, 31 Oct 2023 19:41:56 +0000
Subject: [PATCH 16/18] Return _VizShield

---
 src/viztracer/tracer.py    |  3 ++-
 src/viztracer/viztracer.py | 20 +++++++++++++-------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/viztracer/tracer.py b/src/viztracer/tracer.py
index d3c36909..56ea3a5e 100644
--- a/src/viztracer/tracer.py
+++ b/src/viztracer/tracer.py
@@ -278,7 +278,8 @@ def getts(self) -> float:
         return self._tracer.getts()
 
     def setignorestackcounter(self, value) -> int:
-        return self._tracer.setignorestackcounter(value)
+        # +2 and -2 to compensate for two calls: the current one and the call in the next line
+        return self._tracer.setignorestackcounter(value + 2) - 2
 
     def add_instant(self, name: str, args: Any = None, scope: str = "g") -> None:
         if self.enable:
diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py
index c9b0a6b9..572afe14 100644
--- a/src/viztracer/viztracer.py
+++ b/src/viztracer/viztracer.py
@@ -2,7 +2,6 @@
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
 
 import builtins
-from contextlib import contextmanager
 import multiprocessing
 import os
 import signal
@@ -165,13 +164,20 @@ def log_event(self, event_name: str) -> VizEvent:
         call_frame = sys._getframe(1)
         return VizEvent(self, event_name, call_frame.f_code.co_filename, call_frame.f_lineno)
 
-    @contextmanager
     def shield_ignore(self):
-        prev_ignore_stack = self.setignorestackcounter(0)
-        try:
-            yield
-        finally:
-            self.setignorestackcounter(prev_ignore_stack)
+        class _VizShield:
+            def __init__(self, tracer):
+                self.tracer = tracer
+                self.prev_ignore_stack = None
+
+            def __enter__(self):
+                # 1 is to compensate for the call of the __enter__
+                self.prev_ignore_stack = self.tracer.setignorestackcounter(1)
+
+            def __exit__(self, type, value, trace):
+                self.tracer.setignorestackcounter(self.prev_ignore_stack)
+
+        return _VizShield(self)
 
     def set_afterfork(self, callback: Callable, *args, **kwargs) -> None:
         self._afterfork_cb = callback

From 06abbc3ced2fd96088534fa66c6d9d57d6568337 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Wed, 1 Nov 2023 15:22:18 +0000
Subject: [PATCH 17/18] Update tests for log_sparse

---
 src/viztracer/decorator.py |  3 +--
 src/viztracer/viztracer.py | 19 +++++--------------
 tests/test_logsparse.py    | 33 +++++++++++++++++++++------------
 3 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index b868fb29..2cfcaee1 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -99,8 +99,7 @@ def wrapper(*args, **kwargs) -> Any:
             # The call is made from the module inside, so if `trace_self=False` it will be ignored.
             # To avoid this behavior, we need to reset the counter `ignore_stack_depth`` and then
             # recover it
-            with local_tracer.shield_ignore():
-                return func(*args, **kwargs)
+            return local_tracer.shield_ignore(func, *args, **kwargs)
         else:
             return func(*args, **kwargs)
 
diff --git a/src/viztracer/viztracer.py b/src/viztracer/viztracer.py
index 572afe14..0c41dbdd 100644
--- a/src/viztracer/viztracer.py
+++ b/src/viztracer/viztracer.py
@@ -164,20 +164,11 @@ def log_event(self, event_name: str) -> VizEvent:
         call_frame = sys._getframe(1)
         return VizEvent(self, event_name, call_frame.f_code.co_filename, call_frame.f_lineno)
 
-    def shield_ignore(self):
-        class _VizShield:
-            def __init__(self, tracer):
-                self.tracer = tracer
-                self.prev_ignore_stack = None
-
-            def __enter__(self):
-                # 1 is to compensate for the call of the __enter__
-                self.prev_ignore_stack = self.tracer.setignorestackcounter(1)
-
-            def __exit__(self, type, value, trace):
-                self.tracer.setignorestackcounter(self.prev_ignore_stack)
-
-        return _VizShield(self)
+    def shield_ignore(self, func, *args, **kwargs):
+        prev_ignore_stack = self.setignorestackcounter(0)
+        res = func(*args, **kwargs)
+        self.setignorestackcounter(prev_ignore_stack)
+        return res
 
     def set_afterfork(self, callback: Callable, *args, **kwargs) -> None:
         self._afterfork_cb = callback
diff --git a/tests/test_logsparse.py b/tests/test_logsparse.py
index 66c2a6ab..d85705f2 100644
--- a/tests/test_logsparse.py
+++ b/tests/test_logsparse.py
@@ -1,6 +1,7 @@
 # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
 
+import functools
 import multiprocessing
 import os
 
@@ -151,30 +152,34 @@ def h():
 
 
 class TestLogSparse(CmdlineTmpl):
-    def test_basic(self):
-        def check_func(data):
-            for entry in data["traceEvents"]:
-                self.assertNotEqual(entry["name"], "f")
+    def check_func(self, data, target):
+        names = [entry["name"] for entry in data["traceEvents"]]
+        function_names = [name.split(' ')[0] for name in names if name not in ['process_name', 'thread_name']]
+
+        self.assertEqual(function_names, target)
 
+    def test_basic(self):
         self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"],
                       script=file_basic,
                       expected_output_file="result.json",
-                      expected_entries=1)
+                      expected_entries=1,
+                      check_func=functools.partial(self.check_func, target=['f']))
         self.template(["viztracer", "-o", "result.json", "cmdline_test.py"],
                       script=file_basic,
-                      expected_output_file="result.json",
-                      check_func=check_func)
+                      expected_output_file="result.json")
 
     def test_stack(self):
         self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"],
                       script=file_stack,
                       expected_output_file="result.json",
-                      expected_entries=4)
+                      expected_entries=4,
+                      check_func=functools.partial(self.check_func, target=['f', 'g', 'f', 'g']))
 
         self.template(["viztracer", "-o", "result.json", "--log_sparse", "cmdline_test.py"],
                       script=file_stack_nested,
                       expected_output_file="result.json",
-                      expected_entries=4)
+                      expected_entries=4,
+                      check_func=functools.partial(self.check_func, target=['f', 'g', 'f', 'g']))
 
     def test_without_tracer(self):
         self.template(["python", "cmdline_test.py"], script=file_basic, expected_output_file=None)
@@ -187,6 +192,7 @@ def test_multiprocess(self):
                               script=file_multiprocess,
                               expected_output_file="result.json",
                               expected_entries=3,
+                              check_func=functools.partial(self.check_func, target=['f', 'f', 'f']),
                               concurrency="multiprocessing")
             except Exception as e:
                 # coveragepy has some issue with multiprocess pool
@@ -195,10 +201,13 @@ def test_multiprocess(self):
 
     def test_context_manager(self):
         self.template(["python", "cmdline_test.py"], script=file_context_manager,
-                      expected_output_file="result.json", expected_entries=4)
+                      expected_output_file="result.json", expected_entries=4,
+                      check_func=functools.partial(self.check_func, target=['f', 'g', 'h', 'q']))
 
         self.template(["python", "cmdline_test.py"], script=file_context_manager_logsparse,
-                      expected_output_file="result.json", expected_entries=2)
+                      expected_output_file="result.json", expected_entries=2,
+                      check_func=functools.partial(self.check_func, target=['f', 'q']))
 
         self.template(["python", "cmdline_test.py"], script=file_context_manager_logsparse_stack,
-                      expected_output_file="result.json", expected_entries=2)
+                      expected_output_file="result.json", expected_entries=2,
+                      check_func=functools.partial(self.check_func, target=['g', 'h']))

From 27d578957599a68847a2fa0af9a84846bc069f26 Mon Sep 17 00:00:00 2001
From: alexeykozhevin <alexey.a.kozhevin@gmail.com>
Date: Tue, 7 Nov 2023 15:03:48 +0000
Subject: [PATCH 18/18] Fix linter

---
 src/viztracer/cellmagic.py | 2 +-
 src/viztracer/decorator.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/viztracer/cellmagic.py b/src/viztracer/cellmagic.py
index ab54751f..6c18b7e1 100644
--- a/src/viztracer/cellmagic.py
+++ b/src/viztracer/cellmagic.py
@@ -4,7 +4,7 @@
 try:
     from IPython.core.magic import (Magics, cell_magic,  # type: ignore
                                     magics_class, needs_local_scope)
-    from IPython.core.magic_arguments import (argument, magic_arguments, parse_argstring)
+    from IPython.core.magic_arguments import (argument, magic_arguments, parse_argstring)  # type: ignore
 
     @magics_class
     class VizTracerMagics(Magics):
diff --git a/src/viztracer/decorator.py b/src/viztracer/decorator.py
index 2cfcaee1..6560b8fc 100644
--- a/src/viztracer/decorator.py
+++ b/src/viztracer/decorator.py
@@ -57,7 +57,7 @@ def wrapper(*args, **kwargs) -> Any:
     return inner
 
 
-def _log_sparse_wrapper(func: Optional[Callable], stack_depth: int = 0,
+def _log_sparse_wrapper(func: Callable, stack_depth: int = 0,
                         dynamic_tracer_check: bool = False) -> Callable:
     if not dynamic_tracer_check:
         tracer = get_tracer()