diff --git a/dvc/cli.py b/dvc/cli.py index 2b020d426..34e08a548 100644 --- a/dvc/cli.py +++ b/dvc/cli.py @@ -137,6 +137,14 @@ def get_parent_parser(): """ parent_parser = argparse.ArgumentParser(add_help=False) + parent_parser.add_argument( + "--cprofile", + action="store_true", + default=False, + help=argparse.SUPPRESS, + ) + parent_parser.add_argument("--cprofile-dump", help=argparse.SUPPRESS) + log_level_group = parent_parser.add_mutually_exclusive_group() log_level_group.add_argument( "-q", "--quiet", action="count", default=0, help="Be quiet." diff --git a/dvc/main.py b/dvc/main.py index 07e081253..c02de913d 100644 --- a/dvc/main.py +++ b/dvc/main.py @@ -2,6 +2,7 @@ import errno import logging +from contextlib import contextmanager from dvc import analytics from dvc.cli import parse_args @@ -21,6 +22,25 @@ from dvc.utils import format_link logger = logging.getLogger("dvc") +@contextmanager +def profile(enable, dump): + if not enable: + yield + return + + import cProfile + + prof = cProfile.Profile() + prof.enable() + + yield + + if not dump: + prof.print_stats(sort="cumtime") + return + prof.dump_stats(dump) + + def main(argv=None): # noqa: C901 """Run dvc CLI command. @@ -50,8 +70,9 @@ def main(argv=None): # noqa: C901 logger.trace(args) - cmd = args.func(args) - ret = cmd.run() + with profile(enable=args.cprofile, dump=args.cprofile_dump): + cmd = args.func(args) + ret = cmd.run() except ConfigError: logger.exception("configuration error") ret = 251