Disable leaks warning if tracing Python allocators

We warn when generating a leaks report when pymalloc is in use, because
the memory held by pymalloc would appear to us as a leak when it has
actually just been set aside for future use.

But, when `--trace-python-allocators` is used, we do have visibility
into whether the memory has set aside for future use, and so we don't
need to show this warning in that case.

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>

src/memray/reporters/templates/base.html

@@ -55,16 +55,21 @@
   <main class="container-fluid">
     <div class="row">
       <div class="col bg-light py-3">
-        {% if show_memory_leaks and metadata.python_allocator == "pymalloc" %}
+        {% if show_memory_leaks and metadata.python_allocator == "pymalloc" and not metadata.trace_python_allocators %}
         <div class="alert alert-warning alert-dismissible fade show" role="alert">
-          <p><strong>Report generated using "--leaks" using pymalloc
-              allocator</strong></p>
-          <p>This report for memory leaks was generated with the
-            pymalloc allocator active. This can show confusing results because
-            the Python allocator doesn't necessarily release memory to
-            the system when Python objects are deallocated and these can still
-            appear as "leaks". If you want to exclude these, you can run your
-            application with the `PYTHONMALLOC=malloc` environment variable set.
+          <p><strong>Report generated using "--leaks" using pymalloc allocator</strong></p>
+          <p>
+            This report for memory leaks was generated with the pymalloc
+            allocator active, but without tracking enabled for calls to the
+            pymalloc allocator. This will show confusing results because the
+            pymalloc allocator will retain memory in memory pools even after
+            the objects that requested that memory are deallocated, and Memray
+            won't be able to distinguish memory set aside for reuse from leaked
+            memory. You should rerun your application with the
+            `PYTHONMALLOC=malloc` environment variable set or pass the
+            `--trace-python-allocators` flag when profiling your application.
+            <a href="https://bloomberg.github.io/memray/python_allocators.html">
+            Click here</a> for more information.
           </p>
           <button type="button" class="close" data-dismiss="alert" aria-label="Close">
             <span aria-hidden="true">&times;</span>

tests/integration/test_main.py

@@ -112,15 +112,25 @@ def _wait_until_process_blocks(pid: int) -> None:
         time.sleep(0.1)
 
 
-def generate_sample_results(tmp_path, code, *, native=False):
+def generate_sample_results(
+    tmp_path,
+    code,
+    *,
+    native=False,
+    trace_python_allocators=False,
+    disable_pymalloc=False,
+):
     results_file = tmp_path / "result.bin"
+    env = os.environ.copy()
+    env["PYTHONMALLOC"] = "malloc" if disable_pymalloc else "pymalloc"
     subprocess.run(
         [
             sys.executable,
             "-m",
             "memray",
             "run",
             *(["--native"] if native else []),
+            *(["--trace-python-allocators"] if trace_python_allocators else []),
             "--output",
             str(results_file),
             str(code),
@@ -129,6 +139,7 @@ def generate_sample_results(tmp_path, code, *, native=False):
         check=True,
         capture_output=True,
         text=True,
+        env=env,
     )
     return results_file, code
 
@@ -750,6 +761,48 @@ def test_split_threads_subcommand(self, tmp_path, simple_test_file):
         assert output_file.exists()
         assert str(source_file) in output_file.read_text()
 
+    @pytest.mark.parametrize("trace_python_allocators", [True, False])
+    @pytest.mark.parametrize("disable_pymalloc", [True, False])
+    def test_leaks_with_pymalloc_warning(
+        self,
+        tmp_path,
+        simple_test_file,
+        trace_python_allocators,
+        disable_pymalloc,
+    ):
+        results_file, _ = generate_sample_results(
+            tmp_path,
+            simple_test_file,
+            native=True,
+            trace_python_allocators=trace_python_allocators,
+            disable_pymalloc=disable_pymalloc,
+        )
+        output_file = tmp_path / "output.html"
+        warning_expected = not trace_python_allocators and not disable_pymalloc
+
+        # WHEN
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "memray",
+                "flamegraph",
+                "--leaks",
+                str(results_file),
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        # THEN
+        output_file = tmp_path / "memray-flamegraph-result.html"
+        assert output_file.exists()
+        assert warning_expected == (
+            'Report generated using "--leaks" using pymalloc allocator'
+            in output_file.read_text()
+        )
+
 
 class TestSummarySubCommand:
     def test_summary_generated(self, tmp_path, simple_test_file):