diff --git a/dvc/_debug.py b/dvc/_debug.py index 917368620..e62c9729b 100644 --- a/dvc/_debug.py +++ b/dvc/_debug.py @@ -1,16 +1,13 @@ from contextlib import ExitStack, contextmanager from datetime import datetime -from typing import TYPE_CHECKING, Callable, Union +from typing import TYPE_CHECKING, Callable, Optional, Union if TYPE_CHECKING: from argparse import Namespace @contextmanager -def viztracer_profile( - path: Union[Callable[[], str], str], - max_stack_depth: int = -1, -): +def viztracer_profile(path: Union[Callable[[], str], str], depth: int = -1): try: import viztracer # pylint: disable=import-error except ImportError: @@ -18,7 +15,7 @@ def viztracer_profile( yield return - tracer = viztracer.VizTracer(max_stack_depth=max_stack_depth) + tracer = viztracer.VizTracer(max_stack_depth=depth) tracer.start() yield @@ -109,6 +106,15 @@ def debug(): raise # prevent from jumping ahead +def _get_path_func(tool: str, ext: str): + fmt = f"{tool}.dvc-{{now:%Y%m%d}}_{{now:%H%M%S}}.{ext}" + + def func(now: Optional["datetime"] = None) -> str: + return fmt.format(now=now or datetime.now()) + + return func + + @contextmanager def debugtools(args: "Namespace" = None, **kwargs): kw = vars(args) if args else {} @@ -122,18 +128,13 @@ def debugtools(args: "Namespace" = None, **kwargs): if kw.get("instrument") or kw.get("instrument_open"): stack.enter_context(instrument(kw.get("instrument_open", False))) if kw.get("yappi"): - output = "callgrind.dvc-{0:%Y%m%d}_{0:%H%M%S}.out" - stack.enter_context( - yappi_profile(path=lambda: output.format(datetime.now())) - ) + path_func = _get_path_func("callgrind", "out") + stack.enter_context(yappi_profile(path=path_func)) if kw.get("viztracer") or kw.get("viztracer_depth"): - output = "viztracer.dvc-{0:%Y%m%d}_{0:%H%M%S}.json" - stack.enter_context( - viztracer_profile( - path=lambda: output.format(datetime.now()), - max_stack_depth=kw.get("viztracer_depth") or -1, - ) - ) + path_func = _get_path_func("viztracer", "json") + depth = kw.get("viztracer_depth") or -1 + prof = viztracer_profile(path=path_func, depth=depth) + stack.enter_context(prof) yield diff --git a/tests/unit/command/test_debug.py b/tests/unit/command/test_debug.py deleted file mode 100644 index b97c4499b..000000000 --- a/tests/unit/command/test_debug.py +++ /dev/null @@ -1,19 +0,0 @@ -from dvc.cli import main - - -def test_viztracer(tmp_dir, dvc, mocker): - viztracer_profile = mocker.patch("dvc._debug.viztracer_profile") - - assert main(["status", "--viztracer"]) == 0 - - args = viztracer_profile.call_args[1] - assert callable(args["path"]) - assert args["max_stack_depth"] == -1 - - assert main(["status", "--viztracer", "--viztracer-depth", "5"]) == 0 - args = viztracer_profile.call_args[1] - assert args["max_stack_depth"] == 5 - - assert main(["status", "--viztracer-depth", "2"]) == 0 - args = viztracer_profile.call_args[1] - assert args["max_stack_depth"] == 2