viztracer 1.0.4__cp310-cp310-win_amd64.whl → 1.1.0__cp310-cp310-win_amd64.whl

viztracer/__init__.py

@@ -1,7 +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
 
-__version__ = "1.0.4"
+__version__ = "1.1.0"
 
 
 from .cellmagic import load_ipython_extension

viztracer/code_monkey.py

@@ -5,7 +5,7 @@ import ast
 import copy
 import re
 from functools import reduce
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable
 
 from .util import color_print
 
@@ -18,13 +18,13 @@ class AstTransformer(ast.NodeTransformer):
         self.curr_lineno: int = 0
         self.log_func_exec_enable: bool = False
 
-    def visit_Assign(self, node: ast.Assign) -> Union[ast.stmt, list[ast.stmt]]:
+    def visit_Assign(self, node: ast.Assign) -> ast.stmt | list[ast.stmt]:
         return self._visit_generic_assign(node)
 
-    def visit_AugAssign(self, node: ast.AugAssign) -> Union[ast.stmt, list[ast.stmt]]:
+    def visit_AugAssign(self, node: ast.AugAssign) -> ast.stmt | list[ast.stmt]:
         return self._visit_generic_assign(node)
 
-    def visit_AnnAssign(self, node: ast.AnnAssign) -> Union[ast.stmt, list[ast.stmt]]:
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.stmt | list[ast.stmt]:
         return self._visit_generic_assign(node)
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
@@ -69,13 +69,13 @@ class AstTransformer(ast.NodeTransformer):
                 node.body = instrumented_nodes + node.body
         return node
 
-    def visit_Raise(self, node: ast.Raise) -> Union[ast.AST, list[ast.AST]]:
+    def visit_Raise(self, node: ast.Raise) -> ast.AST | list[ast.AST]:
         if self.inst_type == "log_exception":
             instrument_node = self.get_instrument_node_by_node("Exception", node.exc)
             return [instrument_node, node]
         return node
 
-    def _visit_generic_assign(self, node: Union[ast.Assign, ast.AugAssign, ast.AnnAssign]) -> list[ast.stmt]:
+    def _visit_generic_assign(self, node: ast.Assign | ast.AugAssign | ast.AnnAssign) -> list[ast.stmt]:
         self.generic_visit(node)
         ret: list[ast.stmt] = [node]
         self.curr_lineno = node.lineno
@@ -174,7 +174,7 @@ class AstTransformer(ast.NodeTransformer):
         else:
             raise ValueError(f"{name} is not supported")
 
-    def get_instrument_node_by_node(self, trigger: str, node: Optional[ast.expr]) -> ast.Expr:
+    def get_instrument_node_by_node(self, trigger: str, node: ast.expr | None) -> ast.Expr:
         var_node: ast.expr
         if node is None:
             name = f"{trigger}"
@@ -235,17 +235,14 @@ class AstTransformer(ast.NodeTransformer):
                 n.ctx = ast.Load()  # type: ignore
         return new_node
 
-    def get_string_of_expr(self, node: Union[ast.expr, ast.slice]) -> str:
+    def get_string_of_expr(self, node: ast.expr | ast.slice) -> str:
         """
         Try to do "unparse" of the node
         """
         if isinstance(node, ast.Name):
             return node.id
         elif isinstance(node, ast.Constant):
-            if isinstance(node.value, str):
-                return f"'{node.value}'"
-            else:
-                return f"{node.value}"
+            return repr(node.value)
         elif isinstance(node, ast.Attribute):
             return f"{self.get_string_of_expr(node.value)}.{node.attr}"
         elif isinstance(node, ast.Subscript):
@@ -333,7 +330,7 @@ class CodeMonkey:
     def __init__(self, file_name: str) -> None:
         self.file_name: str = file_name
         self._compile: Callable = compile
-        self.source_processor: Optional[SourceProcessor] = None
+        self.source_processor: SourceProcessor | None = None
         self.ast_transformers: list[AstTransformer] = []
 
     def add_instrument(self, inst_type: str, inst_args: dict[str, dict]) -> None:

viztracer/decorator.py

@@ -5,7 +5,7 @@ import functools
 import multiprocessing
 import os
 import time
-from typing import Any, Callable, Optional, TypeVar, Union, overload
+from typing import Any, Callable, TypeVar, overload
 
 from .viztracer import VizTracer, get_tracer
 
@@ -15,19 +15,18 @@ R = TypeVar("R")
 
 @overload
 def ignore_function(method: None,
-                    tracer: Optional[VizTracer] = None) -> Callable[[Callable[..., R]], Callable[..., R]]:
+                    tracer: VizTracer | None = None) -> Callable[[Callable[..., R]], Callable[..., R]]:
     pass  # pragma: no cover
 
 
 @overload
 def ignore_function(method: Callable[..., R],
-                    tracer: Optional[VizTracer] = None) -> Callable[..., R]:
+                    tracer: VizTracer | None = None) -> Callable[..., R]:
     pass  # pragma: no cover
 
 
-def ignore_function(method: Optional[Callable[..., R]] = None,
-                    tracer: Optional[VizTracer] = None) -> Union[Callable[..., R],
-                                                                 Callable[[Callable[..., R]], Callable[..., R]]]:
+def ignore_function(method: Callable[..., R] | None = None,
+                    tracer: VizTracer | None = None) -> Callable[..., R] | Callable[[Callable[..., R]], Callable[..., R]]:
 
     def inner(func: Callable[..., R]) -> Callable[..., R]:
 
@@ -65,10 +64,9 @@ def trace_and_save(method: Callable[..., R],
     pass  # pragma: no cover
 
 
-def trace_and_save(method: Optional[Callable[..., R]] = None,
+def trace_and_save(method: Callable[..., R] | None = None,
                    output_dir: str = "./",
-                   **viztracer_kwargs) -> Union[Callable[..., R],
-                                                Callable[[Callable[..., R]], Callable[..., R]]]:
+                   **viztracer_kwargs) -> Callable[..., R] | Callable[[Callable[..., R]], Callable[..., R]]:
 
     def inner(func: Callable[..., R]) -> Callable[..., R]:
 
@@ -158,10 +156,9 @@ def log_sparse(func: Callable[..., R],
     pass  # pragma: no cover
 
 
-def log_sparse(func: Optional[Callable[..., R]] = None,
+def log_sparse(func: Callable[..., R] | None = None,
                stack_depth: int = 0,
-               dynamic_tracer_check: bool = False) -> Union[Callable[..., R],
-                                                            Callable[[Callable[..., R]], Callable[..., R]]]:
+               dynamic_tracer_check: bool = False) -> Callable[..., R] | Callable[[Callable[..., R]], Callable[..., R]]:
     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)

viztracer/event_base.py

@@ -2,7 +2,7 @@
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
 
 import functools
-from typing import Any, Callable, Literal, Optional
+from typing import Any, Callable, Literal
 
 from .viztracer import VizTracer
 
@@ -59,7 +59,7 @@ class _EventBase:
         self._viztracer_log()
 
     @staticmethod
-    def triggerlog(method: Optional[Callable] = None,
+    def triggerlog(method: Callable | None = None,
                    when: Literal["after", "before", "both"] = "after") -> Callable:
         if when not in ["after", "before", "both"]:
             raise ValueError(f"when has to be one of 'after', 'before' or 'both', not {when}")

viztracer/functree.py

@@ -4,18 +4,18 @@
 import bisect
 import copy
 import re
-from typing import Any, Generator, Optional
+from typing import Any, Generator
 
 
 class FuncTreeNode:
     name_regex = r"(.*) \((.*?):([0-9]+)\)"
 
-    def __init__(self, event: Optional[dict[str, Any]] = None) -> None:
-        self.filename: Optional[str] = None
-        self.lineno: Optional[int] = None
-        self.is_python: Optional[bool] = False
-        self.funcname: Optional[str] = None
-        self.parent: Optional[FuncTreeNode] = None
+    def __init__(self, event: dict[str, Any] | None = None) -> None:
+        self.filename: str | None = None
+        self.lineno: int | None = None
+        self.is_python: bool | None = False
+        self.funcname: str | None = None
+        self.parent: FuncTreeNode | None = None
         self.children: list[FuncTreeNode] = []
         self.start: float = - (2 ** 64)
         self.end: float = 2 ** 64

viztracer/main.py

@@ -19,7 +19,7 @@ import time
 import types
 import re
 from types import CodeType
-from typing import Any, Optional, Union
+from typing import Any
 
 from . import __version__
 from .code_monkey import CodeMonkey
@@ -31,12 +31,12 @@ from .viztracer import VizTracer
 # For all the procedures in VizUI, return a tuple as the result
 # The first element bool indicates whether the procedure succeeds
 # The second element is the error message if it fails.
-VizProcedureResult = tuple[bool, Optional[str]]
+VizProcedureResult = tuple[bool, str | None]
 
 
 class VizUI:
     def __init__(self) -> None:
-        self.tracer: Optional[VizTracer] = None
+        self.tracer: VizTracer | None = None
         self.parser: argparse.ArgumentParser = self.create_parser()
         self.verbose: int = 1
         self.ofile: str = "result.json"
@@ -199,7 +199,7 @@ class VizUI:
     def parse(self, argv: list[str]) -> VizProcedureResult:
         # If -- or --run exists, all the commands after --/--run are the commands we need to run
         # We need to filter those out, they might conflict with our arguments
-        idx: Optional[int] = None
+        idx: int | None = None
         if "--" in argv[1:]:
             idx = argv.index("--")
         elif "--run" in argv[1:]:
@@ -295,7 +295,7 @@ class VizUI:
 
         return True, None
 
-    def search_file(self, file_name: str) -> Optional[str]:
+    def search_file(self, file_name: str) -> str | None:
         if os.path.isfile(file_name):
             return file_name
 
@@ -326,7 +326,7 @@ class VizUI:
             return self.uninstall()
         elif self.options.cmd_string is not None:
             return self.run_string()
-        elif self.options.module:
+        elif self.options.module is not None:
             return self.run_module()
         elif self.command:
             return self.run_command()
@@ -342,7 +342,7 @@ class VizUI:
             self.parser.print_help()
             return True, None
 
-    def run_code(self, code: Union[CodeType, str], global_dict: dict[str, Any]) -> VizProcedureResult:
+    def run_code(self, code: CodeType | str, global_dict: dict[str, Any]) -> VizProcedureResult:
         options = self.options
         self.parent_pid = os.getpid()
 
@@ -386,16 +386,12 @@ class VizUI:
         if not options.log_exit:
             tracer.stop(stop_option="flush_as_finish")
 
-            # Clear to global_dict to release all references.
-            # This is helpful for some deadlock issues.
-            global_dict.clear()
-
         # issue141 - concurrent.future requires a proper release by executing
         # threading._threading_atexits or it will deadlock if not explicitly
         # release the resource in the code
         # Python 3.9+ has this issue
         if threading._threading_atexits:  # type: ignore
-            for atexit_call in threading._threading_atexits:  # type: ignore
+            for atexit_call in reversed(threading._threading_atexits):  # type: ignore
                 atexit_call()
             threading._threading_atexits = []  # type: ignore
 
@@ -526,7 +522,7 @@ class VizUI:
         print(__version__)
         return True, None
 
-    def _check_attach_availability(self) -> tuple[bool, Optional[str]]:
+    def _check_attach_availability(self) -> tuple[bool, str | None]:
         if sys.platform == "win32":
             return False, "VizTracer does not support this feature on Windows"
 

viztracer/modules/snaptrace.c

@@ -164,9 +164,11 @@ verbose_printf(TracerObject* self, int v, const char* fmt, ...)
 void
 clear_stack(struct FunctionNode** stack_top) {
     Py_CLEAR((*stack_top)->args);
+    Py_CLEAR((*stack_top)->func);
     while ((*stack_top)->prev) {
         (*stack_top) = (*stack_top) -> prev;
         Py_CLEAR((*stack_top)->args);
+        Py_CLEAR((*stack_top)->func);
     }
 }
 
@@ -419,7 +421,7 @@ tracer_pycall_callback(TracerObject* self, PyCodeObject* code)
         info->paused = 1;
         for (size_t i = 0; i < sizeof(curr_task_getters)/sizeof(curr_task_getters[0]); i++) {
             if (curr_task_getters[i] != NULL) {
-                curr_task = PyObject_CallObject(curr_task_getters[i], NULL);
+                curr_task = PyObject_CallNoArgs(curr_task_getters[i]);
                 if (!curr_task) {
                     PyErr_Clear();  // RuntimeError, probably
                     curr_task = Py_None;
@@ -1202,7 +1204,7 @@ tracer_load(TracerObject* self, PyObject* Py_UNUSED(unused))
                 perror("Failed to access multiprocessing.current_process()");
                 exit(-1);
             }
-            PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
+            PyObject* current_process = PyObject_CallNoArgs(current_process_method);
             if (!current_process_method) {
                 perror("Failed to access multiprocessing.current_process()");
                 exit(-1);
@@ -1270,7 +1272,7 @@ tracer_load(TracerObject* self, PyObject* Py_UNUSED(unused))
                     PyObject* task_name = NULL;
                     if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "get_name")) {
                         PyObject* task_name_method = PyObject_GetAttrString(curr->data.fee.asyncio_task, "get_name");
-                        task_name = PyObject_CallObject(task_name_method, NULL);
+                        task_name = PyObject_CallNoArgs(task_name_method);
                         Py_DECREF(task_name_method);
                     } else if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "name")) {
                         task_name = PyObject_GetAttrString(curr->data.fee.asyncio_task, "name");
@@ -1456,7 +1458,7 @@ tracer_dump(TracerObject* self, PyObject* args, PyObject* kw)
                 perror("Failed to access multiprocessing.current_process()");
                 exit(-1);
             }
-            PyObject* current_process = PyObject_CallObject(current_process_method, NULL);
+            PyObject* current_process = PyObject_CallNoArgs(current_process_method);
             if (!current_process_method) {
                 perror("Failed to access multiprocessing.current_process()");
                 exit(-1);
@@ -1502,7 +1504,7 @@ tracer_dump(TracerObject* self, PyObject* args, PyObject* kw)
                     PyObject* task_name = NULL;
                     if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "get_name")) {
                         PyObject* task_name_method = PyObject_GetAttrString(curr->data.fee.asyncio_task, "get_name");
-                        task_name = PyObject_CallObject(task_name_method, NULL);
+                        task_name = PyObject_CallNoArgs(task_name_method);
                         Py_DECREF(task_name_method);
                     } else if (PyObject_HasAttrString(curr->data.fee.asyncio_task, "name")) {
                         task_name = PyObject_GetAttrString(curr->data.fee.asyncio_task, "name");
@@ -1686,7 +1688,7 @@ tracer_getbasetime(TracerObject* self, PyObject* Py_UNUSED(unused))
 }
 
 static PyObject*
-tracer_setcurrstack(TracerObject* self, PyObject* stack_depth)
+tracer_resetstack(TracerObject* self, PyObject* Py_UNUSED(unused))
 {
     struct ThreadInfo* info = get_thread_info(self);
     if (!info) {
@@ -1694,12 +1696,12 @@ tracer_setcurrstack(TracerObject* self, PyObject* stack_depth)
         return NULL;
     }
 
-    if (!PyLong_Check(stack_depth)) {
-        PyErr_SetString(PyExc_TypeError, "stack_depth must be an integer");
-        return NULL;
-    }
+    info->curr_stack_depth = 0;
+    info->ignore_stack_depth = 0;
 
-    info->curr_stack_depth = PyLong_AsLong(stack_depth);
+    struct FunctionNode* stack_top = info->stack_top;
+    clear_stack(&stack_top);
+    info->stack_top = stack_top;
 
     Py_RETURN_NONE;
 }
@@ -1978,7 +1980,7 @@ static PyMethodDef Tracer_methods[] = {
     {"get_func_args", (PyCFunction)tracer_getfunctionarg, METH_NOARGS, "get current function arg"},
     {"getts", (PyCFunction)tracer_getts, METH_NOARGS, "get timestamp"},
     {"get_base_time", (PyCFunction)tracer_getbasetime, METH_NOARGS, "get base time in nanoseconds"},
-    {"_set_curr_stack_depth", (PyCFunction)tracer_setcurrstack, METH_O, "set current stack depth"},
+    {"reset_stack", (PyCFunction)tracer_resetstack, METH_NOARGS, "reset stack"},
     {"pause", (PyCFunction)tracer_pause, METH_NOARGS, "pause profiling"},
     {"resume", (PyCFunction)tracer_resume, METH_NOARGS, "resume profiling"},
     {"setignorestackcounter", (PyCFunction)tracer_setignorestackcounter, METH_O, "reset ignore stack depth"},

viztracer/modules/util.c

@@ -20,13 +20,9 @@ void
 fprintjson(FILE* fptr, PyObject* obj)
 {
     PyObject* json_dumps = PyObject_GetAttrString(json_module, "dumps");
-    PyObject* call_args = PyTuple_New(1);
-    PyTuple_SetItem(call_args, 0, obj);
-    Py_INCREF(obj);
-    PyObject* args_str = PyObject_CallObject(json_dumps, call_args);
+    PyObject* args_str = PyObject_CallOneArg(json_dumps, obj);
     fprintf(fptr, "%s", PyUnicode_AsUTF8(args_str));
     Py_DECREF(json_dumps);
-    Py_DECREF(call_args);
     Py_DECREF(args_str);
 }
 

viztracer/modules/vcompressor/vc_dump.c

@@ -188,7 +188,6 @@ PyObject*
 json_dumps_to_bytes(PyObject* json_data)
 {
     PyObject* json_ret      = NULL;
-    PyObject* json_args     = NULL;
     PyObject* bytes_data    = NULL;
     PyObject* dumps_func    = NULL;
 
@@ -202,11 +201,7 @@ json_dumps_to_bytes(PyObject* json_data)
     }
 
     // json dumps json_data
-    json_args = PyTuple_New(1);
-    PyTuple_SetItem(json_args, 0, json_data);
-    Py_INCREF(json_data); 
-    json_ret = PyObject_CallObject(dumps_func, json_args);
-    Py_DECREF(json_args);
+    json_ret = PyObject_CallOneArg(dumps_func, json_data);
     if (!json_ret) {
         goto clean_exit;
     }
@@ -240,7 +235,6 @@ json_loads_from_bytes(PyObject* bytes_data)
 {
     PyObject* loads_func    = NULL;
     PyObject* string_data   = NULL;
-    PyObject* json_args     = NULL;
     PyObject* json_data     = NULL;
 
     if (!PyBytes_Check(bytes_data)) {
@@ -261,10 +255,7 @@ json_loads_from_bytes(PyObject* bytes_data)
     }
 
     // convert string to json
-    json_args = PyTuple_New(1);
-    PyTuple_SetItem(json_args, 0, string_data);
-    json_data = PyObject_CallObject(loads_func, json_args);
-    Py_DECREF(json_args);
+    json_data = PyObject_CallOneArg(loads_func, string_data);
     if (!json_data) {
         goto clean_exit;
     }
@@ -283,7 +274,6 @@ clean_exit:
 PyObject*
 compress_bytes(PyObject* bytes_data)
 {
-    PyObject* zlib_args         = NULL;
     PyObject* compress_func     = NULL;
     PyObject* compressed_data   = NULL;
 
@@ -298,12 +288,7 @@ compress_bytes(PyObject* bytes_data)
         goto clean_exit;
     }
 
-    zlib_args = PyTuple_New(1);
-    PyTuple_SetItem(zlib_args, 0, bytes_data);
-    Py_INCREF(bytes_data);
-    compressed_data = PyObject_CallObject(compress_func, zlib_args);
-    // zlib_args steals bytes_data, so release zlib_args will release bytes_data
-    Py_DECREF(zlib_args);
+    compressed_data = PyObject_CallOneArg(compress_func, bytes_data);
     if (!compressed_data) {
         goto clean_exit;
     }
@@ -328,7 +313,6 @@ PyObject*
 decompress_bytes(PyObject* bytes_data)
 {
     // decompress data
-    PyObject* zlib_args = NULL;
     PyObject* decompressed_data = NULL;
     PyObject* decompress_func = NULL;
 
@@ -343,11 +327,7 @@ decompress_bytes(PyObject* bytes_data)
         goto clean_exit;
     }
 
-    zlib_args = PyTuple_New(1);
-    PyTuple_SetItem(zlib_args, 0, bytes_data);
-    Py_INCREF(bytes_data);
-    decompressed_data = PyObject_CallObject(decompress_func, zlib_args);
-    Py_DECREF(zlib_args);
+    decompressed_data = PyObject_CallOneArg(decompress_func, bytes_data);
     if (!decompressed_data) {
         goto clean_exit;
     }

viztracer/modules/vcompressor/vcompressor.c

@@ -95,17 +95,17 @@ parse_trace_events(PyObject* trace_events)
                 Py_INCREF(pid);
                 Py_INCREF(tid);
                 Py_INCREF(name);
-                PyTuple_SetItem(key, 0, pid);
-                PyTuple_SetItem(key, 1, tid);
-                PyTuple_SetItem(key, 2, name);
+                PyTuple_SET_ITEM(key, 0, pid);
+                PyTuple_SET_ITEM(key, 1, tid);
+                PyTuple_SET_ITEM(key, 2, name);
 
                 // This indicates that if there's args in fee
                 fee_args = PyDict_GetItemString(event, "args");
                 if (fee_args) {
-                    PyTuple_SetItem(key, 3, Py_True);
+                    PyTuple_SET_ITEM(key, 3, Py_True);
                     Py_INCREF(Py_True);
                 } else {
-                    PyTuple_SetItem(key, 3, Py_False);
+                    PyTuple_SET_ITEM(key, 3, Py_False);
                     Py_INCREF(Py_False);
                 }
 
@@ -126,11 +126,11 @@ parse_trace_events(PyObject* trace_events)
                 
                 Py_INCREF(ts);
                 Py_INCREF(dur);
-                PyTuple_SetItem(ts_dur_tuple, 0, ts);
-                PyTuple_SetItem(ts_dur_tuple, 1, dur);
+                PyTuple_SET_ITEM(ts_dur_tuple, 0, ts);
+                PyTuple_SET_ITEM(ts_dur_tuple, 1, dur);
 
                 if (fee_args) {
-                    PyTuple_SetItem(ts_dur_tuple, 2, fee_args);
+                    PyTuple_SET_ITEM(ts_dur_tuple, 2, fee_args);
                     Py_INCREF(fee_args);
                 }
                 
@@ -154,8 +154,8 @@ parse_trace_events(PyObject* trace_events)
                 // PyTuple_SetItem steals reference
                 Py_INCREF(pid);
                 Py_INCREF(tid);
-                PyTuple_SetItem(id_key, 0, pid);
-                PyTuple_SetItem(id_key, 1, tid);
+                PyTuple_SET_ITEM(id_key, 0, pid);
+                PyTuple_SET_ITEM(id_key, 1, tid);
 
                 if (PyUnicode_CompareWithASCIIString(name, "process_name") == 0) {
                     PyDict_SetItem(process_names, id_key, args_name);
@@ -185,9 +185,9 @@ parse_trace_events(PyObject* trace_events)
                 Py_INCREF(name);
                 Py_INCREF(pid);
                 Py_INCREF(tid);
-                PyTuple_SetItem(counter_id_key, 0, pid);
-                PyTuple_SetItem(counter_id_key, 1, tid);
-                PyTuple_SetItem(counter_id_key, 2, name);
+                PyTuple_SET_ITEM(counter_id_key, 0, pid);
+                PyTuple_SET_ITEM(counter_id_key, 1, tid);
+                PyTuple_SET_ITEM(counter_id_key, 2, name);
                 if (!PyDict_Contains(counter_events, counter_id_key)){
                     counter_event_dict = PyDict_New();
                     PyDict_SetItem(counter_events, counter_id_key, counter_event_dict);

viztracer/patch.py

@@ -10,7 +10,7 @@ import shutil
 import sys
 import textwrap
 from multiprocessing import Process
-from typing import Any, Callable, Sequence, Union, no_type_check
+from typing import Any, Callable, Sequence, no_type_check
 
 from .viztracer import VizTracer
 
@@ -28,7 +28,7 @@ def patch_subprocess(viz_args: list[str]) -> None:
 
     def build_command(args: Sequence[str]) -> list[str] | None:
         py_args: list[str] = []
-        mode = []
+        mode: list[str] | None = []
         script = None
         args_iter = iter(args[1:])
         for arg in args_iter:
@@ -51,7 +51,11 @@ def patch_subprocess(viz_args: list[str]) -> None:
                         py_args.append(f"-{cm_py_args}")
                     mode = [f"-{cm}"]
                     # -m mod | -mmod
-                    mode.append(cm_arg or next(args_iter, None))
+                    cm_arg = cm_arg or next(args_iter, None)
+                    if cm_arg is not None:
+                        mode.append(cm_arg)
+                    else:
+                        mode = None
                 break
 
             # -pyopts
@@ -59,7 +63,7 @@ def patch_subprocess(viz_args: list[str]) -> None:
 
         if script:
             return [sys.executable, *py_args, "-m", "viztracer", "--quiet", *viz_args, "--", script, *args_iter]
-        elif mode and mode[-1] is not None:
+        elif mode:
             return [sys.executable, *py_args, "-m", "viztracer", "--quiet", *viz_args, *mode, "--", *args_iter]
         return None
 
@@ -78,7 +82,7 @@ def patch_subprocess(viz_args: list[str]) -> None:
         return False
 
     @functools.wraps(subprocess.Popen.__init__)
-    def subprocess_init(self: subprocess.Popen[Any], args: Union[str, Sequence[Any], Any], **kwargs: Any) -> None:
+    def subprocess_init(self: subprocess.Popen[Any], args: str | Sequence[Any] | Any, **kwargs: Any) -> None:
         new_args = args
         if isinstance(new_args, str):
             new_args = shlex.split(new_args, posix=sys.platform != "win32")
@@ -118,7 +122,7 @@ def patch_multiprocessing(tracer: VizTracer, viz_args: list[str]) -> None:
         tracer.register_exit()
 
         tracer.clear()
-        tracer._set_curr_stack_depth(1)
+        tracer.reset_stack()
 
         if tracer._afterfork_cb:
             tracer._afterfork_cb(tracer, *tracer._afterfork_args, **tracer._afterfork_kwargs)

viztracer/report_builder.py

@@ -12,13 +12,13 @@ import os
 import re
 import tokenize
 from string import Template
-from typing import Any, Optional, Sequence, TextIO, Union
+from typing import Any, Sequence, TextIO
 
 from . import __version__
 from .util import color_print, same_line_print
 
 
-def get_json(data: Union[dict, str, tuple[str, dict]]) -> dict[str, Any]:
+def get_json(data: dict[str, Any] | str | tuple[str, dict]) -> dict[str, Any]:
     # This function will return a json object if data is already json object
     # or a opened file or a file path
     if isinstance(data, dict):
@@ -65,11 +65,11 @@ def get_json(data: Union[dict, str, tuple[str, dict]]) -> dict[str, Any]:
 class ReportBuilder:
     def __init__(
             self,
-            data: Union[Sequence[Union[str, dict, tuple[str, dict]]], dict],
+            data: Sequence[str | dict | tuple[str, dict]] | dict[str, Any],
             verbose: int = 1,
             align: bool = False,
             minimize_memory: bool = False,
-            base_time: Optional[int] = None) -> None:
+            base_time: int | None = None) -> None:
         self.data = data
         self.verbose = verbose
         self.combined_json: dict = {}
@@ -145,7 +145,7 @@ class ReportBuilder:
                 self.combined_json["file_info"]["files"].update(one["file_info"]["files"])
                 self.combined_json["file_info"]["functions"].update(one["file_info"]["functions"])
 
-    def align_events(self, original_events: list[dict[str, Any]], sync_marker: Optional[float] = None) -> list[dict[str, Any]]:
+    def align_events(self, original_events: list[dict[str, Any]], sync_marker: float | None = None) -> list[dict[str, Any]]:
         """
         Apply an offset to all the trace events, making the start timestamp 0
         This is useful when comparing multiple runs of the same script
@@ -164,7 +164,7 @@ class ReportBuilder:
                 event["ts"] -= offset_ts
         return original_events
 
-    def prepare_json(self, file_info: bool = True, display_time_unit: Optional[str] = None) -> None:
+    def prepare_json(self, file_info: bool = True, display_time_unit: str | None = None) -> None:
         # This will prepare self.combined_json to be ready to output
         self.load_jsons()
         self.combine_json()
@@ -208,7 +208,7 @@ class ReportBuilder:
                 del func_dict[func]
 
     @classmethod
-    def get_source_from_filename(cls, filename: str) -> Optional[str]:
+    def get_source_from_filename(cls, filename: str) -> str | None:
         if filename.startswith("<frozen "):
             m = re.match(r"<frozen (.*)>", filename)
             if not m:
@@ -258,7 +258,7 @@ class ReportBuilder:
                 else:
                     output_file.write(json.dumps(self.combined_json))  # type: ignore
 
-    def save(self, output_file: Union[str, TextIO] = "result.html", file_info: bool = True) -> None:
+    def save(self, output_file: str | TextIO = "result.html", file_info: bool = True) -> None:
         if isinstance(output_file, str):
             file_type = output_file.split(".")[-1]