Repository: bloomberg/memray (python)

Issue:
Make the `%%memray_flamegraph` IPython magic use aggregated capture files
See #519 for context. The `%%memray_flamegraph` IPython magic predates [aggregated capture files](https://bloomberg.github.io/memray/run.html#aggregated-capture-files), but considering that the capture file produced by `%%memray_flamegraph` is only used for the reports that aggregated capture files support, we should probably use aggregated capture files by default for the `%%memray_flamegraph` magic.

Pull request description:
IPython improvements
Closes #520 
See #519 

Patch:
diff --git a/docs/jupyter_magic.rst b/docs/jupyter_magic.rst
index 81ae52b..2b349a5 100644
--- a/docs/jupyter_magic.rst
+++ b/docs/jupyter_magic.rst
@@ -1,3 +1,5 @@
+.. _Jupyter integration:
+
 Jupyter Integration
 ===================
 
diff --git a/news/538.feature.1.rst b/news/538.feature.1.rst
new file mode 100644
index 0000000..d65015a
--- /dev/null
+++ b/news/538.feature.1.rst
@@ -0,0 +1,1 @@
+Automatically use aggregated capture files for the :ref:`Jupyter magic <Jupyter integration>` whenever possible, reducing the amount of disk space needed for temporary files.
diff --git a/news/538.feature.2.rst b/news/538.feature.2.rst
new file mode 100644
index 0000000..462a02c
--- /dev/null
+++ b/news/538.feature.2.rst
@@ -0,0 +1,1 @@
+Allow ``--temporal`` and ``--max-memory-records`` to be used with our :ref:`Jupyter magic <Jupyter integration>`.
diff --git a/src/memray/_ipython/flamegraph.py b/src/memray/_ipython/flamegraph.py
index 27a34ca..b4037af 100644
--- a/src/memray/_ipython/flamegraph.py
+++ b/src/memray/_ipython/flamegraph.py
@@ -18,12 +18,13 @@ from memray.commands.common import warn_if_not_enough_symbols
 from memray.reporters.flamegraph import FlameGraphReporter
 
 TEMPLATE = """\
-from memray import Tracker
+from memray import Tracker, FileFormat
 with Tracker(
     "{dump_file!s}",
     native_traces={native_traces},
     trace_python_allocators={trace_python_allocators},
     follow_fork={follow_fork},
+    file_format=FileFormat.{file_format},
 ) as tracker:
 {code}
 """
@@ -83,13 +84,21 @@ def argument_parser() -> argparse.ArgumentParser:
         const=1,
     )
 
+    parser.add_argument(
+        "--temporal",
+        help=(
+            "Generate a dynamic flame graph that can analyze"
+            " allocations in a user-selected time range."
+        ),
+        action="store_true",
+        default=False,
+    )
     parser.add_argument(
         "--split-threads",
         help="Do not merge allocations across threads",
         action="store_true",
         default=False,
     )
-
     parser.add_argument(
         "--inverted",
         help=(
@@ -99,6 +108,12 @@ def argument_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--max-memory-records",
+        help="Maximum number of memory records to display",
+        type=int,
+        default=None,
+    )
 
     return parser
 
@@ -118,6 +133,15 @@ class FlamegraphMagics(Magics):
             # It already printed a message, just return control to IPython.
             return
 
+        if options.temporal and options.temporary_allocation_threshold >= 0:
+            raise UsageError(
+                "Can't create a temporal flame graph of temporary allocations"
+            )
+        elif options.temporal or options.temporary_allocation_threshold >= 0:
+            file_format = "ALL_ALLOCATIONS"
+        else:
+            file_format = "AGGREGATED_ALLOCATIONS"
+
         results_dir = Path("memray-results")
         results_dir.mkdir(exist_ok=True)
 
@@ -129,6 +153,7 @@ class FlamegraphMagics(Magics):
             trace_python_allocators=options.trace_python_allocators,
             follow_fork=options.follow_fork,
             code=indent(cell, " " * 4),
+            file_format=file_format,
         )
         self.shell.run_cell(code)
 
@@ -136,30 +161,60 @@ class FlamegraphMagics(Magics):
 
         reporter = None
 
-        with FileReader(dump_file, report_progress=True) as reader:
+        kwargs = {}
+        if options.max_memory_records is not None:
+            kwargs["max_memory_records"] = options.max_memory_records
+
+        with FileReader(dump_file, report_progress=True, **kwargs) as reader:
             if reader.metadata.has_native_traces:
                 warn_if_not_enough_symbols()
 
-            if options.show_memory_leaks:
-                snapshot = reader.get_leaked_allocation_records(
-                    merge_threads=merge_threads
-                )
-            elif options.temporary_allocation_threshold >= 0:
-                snapshot = reader.get_temporary_allocation_records(
-                    threshold=options.temporary_allocation_threshold,
-                    merge_threads=merge_threads,
-                )
+            if options.temporal:
+                if options.show_memory_leaks:
+                    temporal_snapshot = reader.get_temporal_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                    reporter = FlameGraphReporter.from_temporal_snapshot(
+                        temporal_snapshot,
+                        memory_records=tuple(reader.get_memory_snapshots()),
+                        native_traces=reader.metadata.has_native_traces,
+                        high_water_mark_by_snapshot=None,
+                        inverted=options.inverted,
+                    )
+                else:
+                    recs, hwms = reader.get_temporal_high_water_mark_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                    reporter = FlameGraphReporter.from_temporal_snapshot(
+                        recs,
+                        memory_records=tuple(reader.get_memory_snapshots()),
+                        native_traces=reader.metadata.has_native_traces,
+                        high_water_mark_by_snapshot=hwms,
+                        inverted=options.inverted,
+                    )
             else:
-                snapshot = reader.get_high_watermark_allocation_records(
-                    merge_threads=merge_threads
+                if options.show_memory_leaks:
+                    snapshot = reader.get_leaked_allocation_records(
+                        merge_threads=merge_threads
+                    )
+                elif options.temporary_allocation_threshold >= 0:
+                    snapshot = reader.get_temporary_allocation_records(
+                        threshold=options.temporary_allocation_threshold,
+                        merge_threads=merge_threads,
+                    )
+                else:
+                    snapshot = reader.get_high_watermark_allocation_records(
+                        merge_threads=merge_threads
+                    )
+
+                memory_records = tuple(reader.get_memory_snapshots())
+                reporter = FlameGraphReporter.from_snapshot(
+                    snapshot,
+                    memory_records=memory_records,
+                    native_traces=options.native,
+                    inverted=options.inverted,
                 )
-            memory_records = tuple(reader.get_memory_snapshots())
-            reporter = FlameGraphReporter.from_snapshot(
-                snapshot,
-                memory_records=memory_records,
-                native_traces=options.native,
-                inverted=options.inverted,
-            )
+
         assert reporter is not None
         flamegraph_path = Path(tempdir) / "flamegraph.html"
         with open(flamegraph_path, "w") as f:


Test patch:
diff --git a/tests/integration/test_ipython.py b/tests/integration/test_ipython.py
index 5f3caab..11b172c 100644
--- a/tests/integration/test_ipython.py
+++ b/tests/integration/test_ipython.py
@@ -115,13 +115,13 @@ class TestIPython:
         assert "" == stdout
         assert "usage:" in stderr
 
-    def test_passing_valid_argument(self, tmpdir, capsys):
+    def test_passing_valid_arguments(self, tmpdir, capsys):
         # GIVEN
 
         code = [
             "%load_ext memray",
             """
-            %%memray_flamegraph --leaks
+            %%memray_flamegraph --leaks --max-memory-records 5
             x = "a" * 10000
             """,
         ]
@@ -137,3 +137,53 @@ class TestIPython:
         assert "<iframe" in html
         assert "flamegraph.html" in html
         assert "Results saved to" in stdout
+
+    @pytest.mark.parametrize(
+        "args,title",
+        [
+            ("--temporary-allocation-threshold=2", "flamegraph report"),
+            ("--leaks", "flamegraph report (memory leaks)"),
+            ("--temporal", "temporal flamegraph report"),
+            ("--leaks --temporal", "temporal flamegraph report (memory leaks)"),
+        ],
+    )
+    def test_report_title_by_report_type(self, tmpdir, capsys, args, title):
+        # GIVEN
+        code = [
+            "%load_ext memray",
+            f"""
+            %%memray_flamegraph {args}
+            x = "a" * 10000
+            """,
+        ]
+
+        # WHEN
+        html = run_in_ipython_shell(tmpdir, code)
+
+        # THEN
+        assert html is not None
+        stdout, _ = capsys.readouterr()
+        assert "<iframe" in html
+        assert "flamegraph.html" in html
+        assert "Results saved to" in stdout
+
+        title_tag = f"<title>memray - {title}</title>"
+        assert title_tag in next(tmpdir.visit("**/flamegraph.html")).read()
+
+    def test_passing_temporal_and_temporary_allocations(self, tmpdir, capsys):
+        # GIVEN
+        code = [
+            "%load_ext memray",
+            """
+            %%memray_flamegraph --temporal --temporary-allocation-threshold=2
+            """,
+        ]
+
+        # WHEN
+        html = run_in_ipython_shell(tmpdir, code)
+
+        # THEN
+        assert html is None
+        stdout, stderr = capsys.readouterr()
+        assert "" == stdout
+        assert "Can't create a temporal flame graph of temporary allocations" in stderr