gh-159560: Add a command-line interface to tracemalloc

Doc/library/cmdline.rst

@@ -50,6 +50,7 @@ The following modules have a command-line interface.
 * :ref:`timeit <timeit-command-line-interface>`
 * :ref:`tokenize <tokenize-cli>`
 * :ref:`trace <trace-cli>`
+* :ref:`tracemalloc <tracemalloc-cli>`
 * :mod:`turtledemo`
 * :ref:`unittest <unittest-command-line-interface>`
 * :ref:`uuid <uuid-cli>`

Doc/library/tracemalloc.rst

@@ -762,3 +762,71 @@ Traceback
               obj = Object()
             File "test.py", line 12
               tb = tracemalloc.get_object_traceback(f())
+
+
+.. _tracemalloc-cli:
+
+Command-line usage
+------------------
+
+.. versionadded:: next
+
+The :mod:`!tracemalloc` module can be invoked as a script to trace the memory
+allocations of a Python program, without modifying its source code:
+
+.. code-block:: sh
+
+   python -m tracemalloc [-o OUTFILE] [-t KEY_TYPE] [--cumulative] [-n NFRAME] [-l LIMIT] (-m module | script.py) [args ...]
+
+Tracing starts just before the program is run.  When the program exits,
+a snapshot is taken and the top statistics of the memory blocks which are
+still allocated are displayed, followed by their total size and the peak
+size of traced memory blocks.  Memory blocks freed before the program exits
+are only reflected in the peak size.  Example of output::
+
+   Top 10 lineno allocations still alive at exit:
+   /home/user/script.py:4: size=977 KiB, count=100, average=10000 B
+   <frozen importlib._bootstrap_external>:511: size=64.7 KiB, count=588, average=113 B
+   ...
+   164 other: size=145 KiB, count=1062
+   Total size: 1411 KiB
+   Peak size: 1412 KiB
+
+The following options are accepted:
+
+.. program:: tracemalloc
+
+.. option:: -h, --help
+
+   Show the help message and exit.
+
+.. option:: -o <outfile>, --outfile <outfile>
+
+   Dump the snapshot to *outfile* instead of displaying statistics.  The
+   snapshot can be loaded later using :meth:`Snapshot.load`.
+
+.. option:: -t <key_type>, --key-type <key_type>
+
+   Group statistics by *key_type*: ``'filename'``, ``'lineno'`` (default) or
+   ``'traceback'``.  See :meth:`Snapshot.statistics`.
+
+.. option:: --cumulative
+
+   Cumulate the size and count of memory blocks of all frames of the
+   traceback of a trace, not only the most recent frame.  This option
+   cannot be used with the ``'traceback'`` key type.
+
+.. option:: -n <nframe>, --nframe <nframe>
+
+   Store at most *nframe* frames in the traceback of a trace (default: 1).
+   See :func:`start`.  Tracing is restarted with this limit if it was already
+   enabled by the :envvar:`PYTHONTRACEMALLOC` environment variable or the
+   ``-X tracemalloc`` command line option.
+
+.. option:: -l <limit>, --limit <limit>
+
+   Display at most *limit* statistics (default: 10).
+
+.. option:: -m <module>
+
+   Run *module* as a script, like ``python -m``.

Doc/whatsnew/3.16.rst

@@ -770,6 +770,15 @@ tkinter
   (Contributed by Serhiy Storchaka in :gh:`153513`.)
 
 
+tracemalloc
+-----------
+
+* Added a :ref:`command-line interface <tracemalloc-cli>` to trace
+  the memory allocations of a script or module without modifying it:
+  ``python -m tracemalloc script.py``.
+  (Contributed by Vadim Smal in :gh:`159560`.)
+
+
 turtle
 ------
 

Lib/test/test_tracemalloc.py

@@ -8,7 +8,8 @@
 from test.support.script_helper import (assert_python_ok, assert_python_failure,
                                         interpreter_requires_environment)
 from test import support
-from test.support import force_not_colorized, warnings_helper
+from test.support import (force_not_colorized,
+                          force_not_colorized_test_class, warnings_helper)
 from test.support import os_helper
 from test.support import threading_helper
 
@@ -1000,6 +1001,149 @@ def test_pymem_alloc0(self):
         assert_python_ok('-X', 'tracemalloc', '-c', code)
 
 
+@force_not_colorized_test_class
+class TestModuleCommandLine(unittest.TestCase):
+    script = textwrap.dedent("""
+        import sys
+        def alloc():
+            return [bytes(10_000) for _ in range(100)]
+        data = alloc()
+        print(f'name={__name__} argv={sys.argv[1:]}')
+    """)
+
+    def setUp(self):
+        self.tmpdir = self.enterContext(os_helper.temp_dir())
+        self.script_path = self.write_script('script.py', self.script)
+
+    def write_script(self, name, source):
+        path = os.path.join(self.tmpdir, name)
+        with open(path, 'w', encoding='utf-8') as fp:
+            fp.write(source)
+        return path
+
+    def run_cli(self, *args, **kwargs):
+        res = assert_python_ok('-m', 'tracemalloc', *args, **kwargs)
+        return res.out.decode()
+
+    def run_cli_failure(self, *args):
+        res = assert_python_failure('-m', 'tracemalloc', *args)
+        return res.err.decode()
+
+    def test_script(self):
+        out = self.run_cli(self.script_path, 'a', '-b', '--c')
+        self.assertIn("name=__main__ argv=['a', '-b', '--c']", out)
+        self.assertIn('Top 10 lineno allocations still alive at exit:', out)
+        self.assertRegex(out, r'script\.py:4: size=\d+ KiB, count=\d+')
+        self.assertIn('Total size: ', out)
+        self.assertIn('Peak size: ', out)
+
+    def test_globals_alive(self):
+        # Memory referenced by globals is alive at exit, even without a
+        # reference cycle between a function and the globals
+        script = self.write_script(
+            'nofunc.py', 'data = [bytes(10_000) for _ in range(100)]\n')
+        out = self.run_cli(script)
+        self.assertRegex(out, r'nofunc\.py:1: size=\d+ KiB, count=\d+')
+
+    def test_module(self):
+        out = self.run_cli('-m', 'script', 'arg', __cwd=self.tmpdir,
+                           __isolated=False)
+        self.assertIn("name=__main__ argv=['arg']", out)
+        self.assertRegex(out, r'script\.py:4: size=\d+ KiB')
+
+    def test_script_sys_path(self):
+        self.write_script('helper.py', 'print("helper imported")\n')
+        script = self.write_script('importer.py', 'import helper\n')
+        out = self.run_cli(script, __isolated=False)
+        self.assertIn('helper imported', out)
+
+    def test_limit(self):
+        out = self.run_cli('-l', '1', self.script_path)
+        lines = out.splitlines()
+        index = lines.index('Top 1 lineno allocations still alive at exit:')
+        self.assertRegex(lines[index + 1], r'script\.py:4: ')
+        self.assertRegex(lines[index + 2], r'^\d+ other: size=')
+
+    def test_key_type_filename(self):
+        out = self.run_cli('-t', 'filename', self.script_path)
+        self.assertRegex(out, r'script\.py:0: size=\d+ KiB')
+
+    def test_key_type_traceback(self):
+        out = self.run_cli('-t', 'traceback', '-n', '5', '-l', '1',
+                           self.script_path)
+        self.assertIn('Top 1 traceback allocations still alive at exit:', out)
+        self.assertIn('    data = alloc()', out)
+        self.assertIn('    return [bytes(10_000) for _ in range(100)]', out)
+        # Frames of the command line interface are not displayed
+        self.assertNotIn('runpy', out)
+        self.assertNotIn('tracemalloc.py', out)
+
+    def test_cumulative(self):
+        out = self.run_cli('--cumulative', '-n', '5', '-l', '2',
+                           self.script_path)
+        self.assertRegex(out, r'script\.py:5: size=\d+ KiB')
+        self.assertRegex(out, r'script\.py:4: size=\d+ KiB')
+        self.assertNotIn('runpy', out)
+        self.assertNotIn(' other: ', out)
+
+    def test_outfile(self):
+        outfile = os.path.join(self.tmpdir, 'snapshot.pickle')
+        out = self.run_cli('-n', '3', '-o', outfile, self.script_path)
+        self.assertNotIn('Total size', out)
+        snapshot = tracemalloc.Snapshot.load(outfile)
+        self.assertEqual(snapshot.traceback_limit, 3)
+        top = snapshot.statistics('lineno')[0]
+        self.assertEqual(top.traceback[0].filename, self.script_path)
+        self.assertEqual(top.traceback[0].lineno, 4)
+
+    def test_already_tracing(self):
+        # -n overrides the number of frames set by -X tracemalloc
+        outfile = os.path.join(self.tmpdir, 'snapshot.pickle')
+        self.run_cli('-n', '3', '-o', outfile, self.script_path,
+                     __isolated=False, PYTHONTRACEMALLOC='1')
+        snapshot = tracemalloc.Snapshot.load(outfile)
+        self.assertEqual(snapshot.traceback_limit, 3)
+
+    def test_exit_code(self):
+        script = self.write_script('exit.py', 'import sys\nsys.exit(3)\n')
+        res = assert_python_failure('-m', 'tracemalloc', script)
+        self.assertEqual(res.rc, 3)
+        self.assertIn(b'Total size: ', res.out)
+
+    def test_exception(self):
+        script = self.write_script('error.py', '1/0\n')
+        res = assert_python_failure('-m', 'tracemalloc', script)
+        self.assertIn(b'ZeroDivisionError', res.err)
+        self.assertIn(b'Total size: ', res.out)
+
+    def test_tracing_stopped(self):
+        # The traced program can stop tracing
+        script = self.write_script(
+            'stop.py', 'import sys, tracemalloc\n'
+                       'tracemalloc.stop()\n'
+                       'sys.exit(3)\n')
+        res = assert_python_failure('-m', 'tracemalloc', script)
+        self.assertEqual(res.rc, 3)
+        self.assertEqual(res.out, b'')
+        self.assertEqual(res.err, b'')
+
+    def test_errors(self):
+        err = self.run_cli_failure()
+        self.assertIn('the following arguments are required: target', err)
+        err = self.run_cli_failure(os.path.join(self.tmpdir, 'missing.py'))
+        self.assertIn('cannot find script', err)
+        err = self.run_cli_failure('-t', 'traceback', '--cumulative',
+                                   self.script_path)
+        self.assertIn("--cumulative cannot be used with key type", err)
+        err = self.run_cli_failure('-l', '0', self.script_path)
+        self.assertIn('--limit must be a positive integer', err)
+        for nframe in INVALID_NFRAME:
+            with self.subTest(nframe=nframe):
+                err = self.run_cli_failure('-n', str(nframe),
+                                           self.script_path)
+                self.assertIn('the number of frames must be in range', err)
+
+
 @unittest.skipIf(_testcapi is None, 'need _testcapi')
 class TestCAPI(unittest.TestCase):
     maxDiff = 80 * 20

Lib/tracemalloc.py

@@ -558,3 +558,138 @@ def take_snapshot():
     traces = _get_traces()
     traceback_limit = get_traceback_limit()
     return Snapshot(traces, traceback_limit)
+
+
+def _main(args=None):
+    import argparse
+    # Import pkgutil now, rather than when runpy imports it, so that its
+    # memory is not attributed to the traced program.
+    import pkgutil
+    import runpy
+    import sys
+
+    parser = argparse.ArgumentParser(
+        prog='python -m tracemalloc',
+        description='Trace memory allocations of a Python script or module '
+                    'and display the top allocations which are still alive '
+                    'when it finishes.',
+    )
+    parser.add_argument('-o', '--outfile',
+                        help='dump the snapshot to OUTFILE instead of '
+                             'displaying statistics; the file can be loaded '
+                             'with Snapshot.load()')
+    parser.add_argument('-t', '--key-type', default='lineno',
+                        choices=('filename', 'lineno', 'traceback'),
+                        help='group statistics by key type '
+                             '(default: %(default)s)')
+    parser.add_argument('--cumulative', action='store_true',
+                        help='cumulate the size and count of memory blocks '
+                             'of all frames of a traceback, not only the '
+                             'most recent frame')
+    parser.add_argument('-n', '--nframe', type=int, default=1,
+                        help='maximum number of frames stored in a traceback '
+                             '(default: %(default)s)')
+    parser.add_argument('-l', '--limit', type=int, default=10,
+                        help='number of statistics to display '
+                             '(default: %(default)s)')
+    parser.add_argument('-m', dest='module', action='store_true',
+                        help='run target as a module (like python -m)')
+    parser.add_argument('target',
+                        help='script file or module name to trace')
+    parser.add_argument('args', nargs=argparse.REMAINDER,
+                        help='arguments passed to the script or module')
+    opts = parser.parse_args(args)
+
+    if opts.cumulative and opts.key_type == 'traceback':
+        parser.error("--cumulative cannot be used with key type 'traceback'")
+    if opts.limit < 1:
+        parser.error('--limit must be a positive integer')
+
+    # The traced program may chdir, so capture the absolute path now.
+    if opts.outfile is not None:
+        opts.outfile = os.path.abspath(opts.outfile)
+
+    if not opts.module:
+        if not os.path.exists(opts.target):
+            parser.error(f'cannot find script {opts.target!r}')
+        # runpy.run_path() inserts directories and zip files in sys.path
+        if (not sys.flags.safe_path
+                and pkgutil.get_importer(opts.target) is None):
+            sys.path[0] = os.path.dirname(os.path.realpath(opts.target))
+    sys.argv[:] = [opts.target, *opts.args]
+
+    # Restart tracing if it was enabled by -X tracemalloc or
+    # PYTHONTRACEMALLOC, so that the requested number of frames is used and
+    # the report only covers the traced program.
+    stop()
+    try:
+        start(opts.nframe)
+    except ValueError as exc:
+        parser.error(str(exc))
+
+    try:
+        # Keep the globals of the traced program alive, so that the memory
+        # they reference is included in the snapshot.
+        if opts.module:
+            globs = runpy.run_module(opts.target, run_name='__main__',
+                                     alter_sys=True)
+        else:
+            globs = runpy.run_path(opts.target, run_name='__main__')
+    finally:
+        # The traced program can stop tracing
+        if is_tracing():
+            size, peak = get_traced_memory()
+            snapshot = _strip_frames(take_snapshot(), {
+                _main.__code__.co_filename,
+                runpy._run_code.__code__.co_filename,
+            })
+            stop()
+            if opts.outfile is not None:
+                snapshot.dump(opts.outfile)
+            else:
+                _print_top(snapshot, opts.key_type, opts.cumulative,
+                           opts.limit, size, peak)
+
+
+def _strip_frames(snapshot, filenames):
+    # Remove the oldest frames belonging to the given files (the frames of
+    # the command line interface), but keep at least the most recent frame.
+    if snapshot.traceback_limit == 1:
+        return snapshot
+    traces = []
+    for domain, size, frames, total_nframe in snapshot.traces._traces:
+        nframe = len(frames)
+        while nframe > 1 and frames[nframe - 1][0] in filenames:
+            nframe -= 1
+        stripped = len(frames) - nframe
+        traces.append((domain, size, frames[:nframe], total_nframe - stripped))
+    return Snapshot(traces, snapshot.traceback_limit)
+
+
+def _print_top(snapshot, key_type, cumulative, limit, size, peak):
+    stats = snapshot.statistics(key_type, cumulative)
+    print(f'Top {min(limit, len(stats))} {key_type} allocations '
+          f'still alive at exit:')
+    for stat in stats[:limit]:
+        if key_type == 'traceback':
+            print(f'{stat.count} memory blocks: '
+                  f'{_format_size(stat.size, False)}')
+            for line in stat.traceback.format():
+                print(line)
+        else:
+            print(stat)
+    if not cumulative and len(stats) > limit:
+        other = stats[limit:]
+        other_size = sum(stat.size for stat in other)
+        other_count = sum(stat.count for stat in other)
+        print(f'{len(other)} other: size={_format_size(other_size, False)}, '
+              f'count={other_count}')
+    print(f'Total size: {_format_size(size, False)}')
+    print(f'Peak size: {_format_size(peak, False)}')
+
+
+if __name__ == '__main__':
+    # Use the classes of the tracemalloc module rather than the ones of
+    # __main__, so that dumped snapshots can be loaded.
+    import tracemalloc
+    tracemalloc._main()

Misc/NEWS.d/next/Library/2026-10-16-17-25-52.gh-issue-159560.XqZA5X.rst

@@ -0,0 +1,3 @@
+Add a command-line interface to :mod:`tracemalloc`: ``python -m tracemalloc
+script.py`` traces the memory allocations of a script or module and displays
+the top allocations still alive at exit. Patch by Vadim Smal.