Merge pull request #954 from linsword13/profile

Tweaks on profile

Author: douglasjacobsen

lib/ramble/ramble/main.py

@@ -502,6 +502,19 @@ def make_argument_parser(**kwargs):
         action="store",
         help="lines of profile output or 'all' (default: 20)",
     )
+    parser.add_argument(
+        "--profile-restrictions",
+        default="",
+        action="store",
+        help=(
+            "Comma-separated restrictions applied to the cProfiler. "
+            "When specified, it takes precedence over `--lines`. "
+            "Example: `--profile-restrictions 'logger,5'` limits the list to functions with "
+            "'logger' in the path, and then shows the first 5 of them. "
+            "See https://docs.python.org/3/library/profile.html#pstats.Stats.print_stats for more "
+            "details."
+        ),
+    )
     parser.add_argument(
         "-v", "--verbose", action="store_true", help="print additional output during builds"
     )
@@ -714,12 +727,13 @@ def _log_command_output(self, out):
 def _profile_wrapper(command, parser, args, unknown_args):
     import cProfile
 
-    try:
-        nlines = int(args.lines)
-    except ValueError:
-        if args.lines != "all":
-            logger.die(f"Invalid number for --lines: {args.lines}")
-        nlines = -1
+    from ramble.util import conversions
+
+    if args.profile_restrictions:
+        restrictions_raw = args.profile_restrictions.split(",")
+    else:
+        restrictions_raw = [args.lines if args.lines != "all" else -1]
+    restrictions = [conversions.convert_to_number(v) for v in restrictions_raw]
 
     # allow comma-separated list of fields
     sortby = ["time"]
@@ -739,9 +753,9 @@ def _profile_wrapper(command, parser, args, unknown_args):
         pr.disable()
 
         # print out profile stats.
-        stats = pstats.Stats(pr)
+        stats = pstats.Stats(pr, stream=sys.stderr)
         stats.sort_stats(*sortby)
-        stats.print_stats(nlines)
+        stats.print_stats(*restrictions)
 
 
 def print_setup_info(*info):

lib/ramble/ramble/test/util/conversions.py

@@ -34,3 +34,19 @@ def test_list_str_to_list(input, expect):
 )
 def test_canonical_none(input, expect):
     assert conversions.canonical_none(input) == expect
+
+
+@pytest.mark.parametrize(
+    "input,expect",
+    [
+        (None, None),
+        ("foo", "foo"),
+        (1, 1),
+        (-1.2, -1.2),
+        ("-1.2", -1.2),
+        ("-10", -10),
+        ("1.0", 1),
+    ],
+)
+def test_convert_to_number(input, expect):
+    assert conversions.convert_to_number(input) == expect

lib/ramble/ramble/util/conversions.py

@@ -37,3 +37,16 @@ def canonical_none(maybe_none):
     if isinstance(maybe_none, str) and maybe_none.lower() == "none":
         return None
     return maybe_none
+
+
+def convert_to_number(val):
+    "Convert (in order of preference) to int, or float, or simply return itself"
+    if not isinstance(val, str):
+        return val
+    try:
+        f_val = float(val)
+        if f_val.is_integer():
+            return int(f_val)
+        return f_val
+    except ValueError:
+        return val

share/ramble/ramble-completion.bash

@@ -264,7 +264,7 @@ complete -o bashdefault -o default -F _bash_completion_ramble ramble
 _ramble() {
     if $list_options
     then
-        RAMBLE_COMPREPLY="-h --help -H --all-help --color -c --config -C --config-scope -d --debug --disable-passthrough -N --disable-logger -P --disable-progress-bar --timestamp --pdb -w --workspace -D --workspace-dir -W --no-workspace --use-workspace-repo -k --insecure -l --enable-locks -L --disable-locks -m --mock --mock-applications --mock-modifiers --mock-package-managers --mock-workflow-managers --mock-base-applications --mock-base-modifiers --mock-base-package-managers --mock-base-workflow-managers -p --profile --sorted-profile --lines -v --verbose --stacktrace -V --version --print-shell-vars"
+        RAMBLE_COMPREPLY="-h --help -H --all-help --color -c --config -C --config-scope -d --debug --disable-passthrough -N --disable-logger -P --disable-progress-bar --timestamp --pdb -w --workspace -D --workspace-dir -W --no-workspace --use-workspace-repo -k --insecure -l --enable-locks -L --disable-locks -m --mock --mock-applications --mock-modifiers --mock-package-managers --mock-workflow-managers --mock-base-applications --mock-base-modifiers --mock-base-package-managers --mock-base-workflow-managers -p --profile --sorted-profile --lines --profile-restrictions -v --verbose --stacktrace -V --version --print-shell-vars"
     else
         RAMBLE_COMPREPLY="attributes clean commands config debug deployment docs edit help info license list mirror on python repo results software-definitions style unit-test workspace"
     fi