pytest-profiling: feat: support cProfile output filtering via regex

Author: IliaFeldgun

pytest-profiling/README.md

@@ -16,7 +16,7 @@ Install using your favourite package installer:
     # or
     easy_install pytest-profiling
 ```
-    
+
 Enable the fixture explicitly in your tests or conftest.py (not required when using setuptools entry points):
 
 ```python
@@ -31,9 +31,17 @@ Once installed, the plugin provides extra options to pytest:
     $ py.test --help
     ...
       Profiling:
-        --profile           generate profiling information
-        --profile-svg       generate profiling graph (using gprof2dot and dot
-                            -Tsvg)
+        --profile             generate profiling information
+        --profile-svg         generate profiling graph (using gprof2dot and dot -Tsvg)
+        --pstats-dir=PSTATS_DIR
+                              configure the dump directory of profile data files
+        --element-number=ELEMENT_NUMBER
+                              defines how many elements will display in a result
+        --element-regex=ELEMENT_REGEX
+                              filters elements displayed using regex
+        --strip-dirs          configure to show/hide the leading path information from
+                              file names
+
 ```
 
 The ``--profile`` and ``profile-svg`` options can be combined with any other option:

pytest-profiling/pytest_profiling.py

@@ -31,11 +31,12 @@ class Profiling(object):
     dot_cmd = None
     gprof2dot_cmd = None
 
-    def __init__(self, svg, dir=None, element_number=20, stripdirs=False):
+    def __init__(self, svg, dir=None, element_number=20, stripdirs=False, element_regex=None):
         self.svg = svg
         self.dir = 'prof' if dir is None else dir[0]
         self.stripdirs = stripdirs
         self.element_number = element_number
+        self.element_regex = element_regex
         self.profs = []
         self.gprof2dot = os.path.abspath(os.path.join(os.path.dirname(sys.executable), 'gprof2dot'))
         if not os.path.isfile(self.gprof2dot):
@@ -98,7 +99,12 @@ def pytest_terminal_summary(self, terminalreporter):
             stats = pstats.Stats(self.combined, stream=terminalreporter)
             if self.stripdirs:
                 stats.strip_dirs()
-            stats.sort_stats('cumulative').print_stats(self.element_number)
+            if self.element_regex:
+                print_args = self.element_regex, self.element_number
+            else:
+                print_args = (self.element_number,)
+            stats.sort_stats('cumulative')
+            stats.print_stats(*print_args)
         if self.svg_name:
             if not self.exit_code:
                 # 0 - SUCCESS
@@ -144,6 +150,8 @@ def pytest_addoption(parser):
                     help="generate profiling graph (using gprof2dot and dot -Tsvg)")
     group.addoption("--pstats-dir", nargs=1,
                     help="configure the dump directory of profile data files")
+    group.addoption("--profile-element-regex", type=str, default=None,
+                    help="filters elements displayed using regex")
     group.addoption("--element-number", action="store", type=int, default=20,
                     help="defines how many elements will display in a result")
     group.addoption("--strip-dirs", action="store_true",
@@ -157,4 +165,5 @@ def pytest_configure(config):
         config.pluginmanager.register(Profiling(config.getvalue('profile_svg'),
                                                 config.getvalue('pstats_dir'),
                                                 element_number=config.getvalue('element_number'),
-                                                stripdirs=config.getvalue('strip_dirs')))
+                                                stripdirs=config.getvalue('strip_dirs'),
+                                                element_regex=config.getvalue('profile_element_regex')))

pytest-profiling/tests/integration/test_profile_integration.py

@@ -69,3 +69,17 @@ def test_profile_chdir(pytestconfig, virtualenv):
         pytestconfig,
         cd=virtualenv.workspace,
     )
+
+
+def test_profile_limited_output(pytestconfig, virtualenv):
+    element_args = ["--element-number=2", r'--profile-element-regex=.*(test_regex|makedirs).*']
+    # element_args = ["--element-number=100"]
+    print(["-m", "pytest", "--profile", *element_args, "tests/unit/test_regex.py"])
+    output = virtualenv.run_with_coverage(
+        ["-m", "pytest", "--profile", *element_args, "tests/unit/test_regex.py"],
+        pytestconfig,
+        cd=virtualenv.workspace,
+    )
+    assert "test_regex.py:4(test_regex)" in output
+    assert "makedirs" in output
+    assert "chdir" not in output

pytest-profiling/tests/unit/test_profile.py

@@ -101,6 +101,10 @@ def test_adds_options():
     group = parser.getgroup.return_value
     group.addoption.assert_any_call("--profile", action="store_true", help=ANY)
     group.addoption.assert_any_call("--profile-svg", action="store_true", help=ANY)
+    group.addoption.assert_any_call("--pstats-dir", nargs=1, help=ANY)
+    group.addoption.assert_any_call("--profile-element-regex", type=str, default=None, help=ANY)
+    group.addoption.assert_any_call("--element-number", action="store", type=int, default=20, help=ANY)
+    group.addoption.assert_any_call("--strip-dirs", action="store_true", help=ANY)
 
 
 def test_configures():