Skip to content

Commit a1fec02

Browse files
committed
Add command-line interface to tracemalloc
1 parent d3ada45 commit a1fec02

4 files changed

Lines changed: 349 additions & 1 deletion

File tree

‎Doc/library/cmdline.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ The following modules have a command-line interface.
5050
* :ref:`timeit <timeit-command-line-interface>`
5151
* :ref:`tokenize <tokenize-cli>`
5252
* :ref:`trace <trace-cli>`
53+
* :ref:`tracemalloc <tracemalloc-cli>`
5354
* :mod:`turtledemo`
5455
* :ref:`unittest <unittest-command-line-interface>`
5556
* :ref:`uuid <uuid-cli>`

‎Doc/library/tracemalloc.rst‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,3 +762,71 @@ Traceback
762762
obj = Object()
763763
File "test.py", line 12
764764
tb = tracemalloc.get_object_traceback(f())
765+
766+
767+
.. _tracemalloc-cli:
768+
769+
Command-line usage
770+
------------------
771+
772+
.. versionadded:: next
773+
774+
The :mod:`!tracemalloc` module can be invoked as a script to trace the memory
775+
allocations of a Python program, without modifying its source code:
776+
777+
.. code-block:: sh
778+
779+
python -m tracemalloc [-o OUTFILE] [-t KEY_TYPE] [--cumulative] [-n NFRAME] [-l LIMIT] (-m module | script.py) [args ...]
780+
781+
Tracing starts just before the program is run. When the program exits,
782+
a snapshot is taken and the top statistics of the memory blocks which are
783+
still allocated are displayed, followed by their total size and the peak
784+
size of traced memory blocks. Memory blocks freed before the program exits
785+
are only reflected in the peak size. Example of output::
786+
787+
Top 10 lineno allocations still alive at exit:
788+
/home/user/script.py:4: size=977 KiB, count=100, average=10000 B
789+
<frozen importlib._bootstrap_external>:511: size=64.7 KiB, count=588, average=113 B
790+
...
791+
164 other: size=145 KiB, count=1062
792+
Total size: 1411 KiB
793+
Peak size: 1412 KiB
794+
795+
The following options are accepted:
796+
797+
.. program:: tracemalloc
798+
799+
.. option:: -h, --help
800+
801+
Show the help message and exit.
802+
803+
.. option:: -o <outfile>, --outfile <outfile>
804+
805+
Dump the snapshot to *outfile* instead of displaying statistics. The
806+
snapshot can be loaded later using :meth:`Snapshot.load`.
807+
808+
.. option:: -t <key_type>, --key-type <key_type>
809+
810+
Group statistics by *key_type*: ``'filename'``, ``'lineno'`` (default) or
811+
``'traceback'``. See :meth:`Snapshot.statistics`.
812+
813+
.. option:: --cumulative
814+
815+
Cumulate the size and count of memory blocks of all frames of the
816+
traceback of a trace, not only the most recent frame. This option
817+
cannot be used with the ``'traceback'`` key type.
818+
819+
.. option:: -n <nframe>, --nframe <nframe>
820+
821+
Store at most *nframe* frames in the traceback of a trace (default: 1).
822+
See :func:`start`. Tracing is restarted with this limit if it was already
823+
enabled by the :envvar:`PYTHONTRACEMALLOC` environment variable or the
824+
``-X tracemalloc`` command line option.
825+
826+
.. option:: -l <limit>, --limit <limit>
827+
828+
Display at most *limit* statistics (default: 10).
829+
830+
.. option:: -m <module>
831+
832+
Run *module* as a script, like ``python -m``.

‎Lib/test/test_tracemalloc.py‎

Lines changed: 145 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
from test.support.script_helper import (assert_python_ok, assert_python_failure,
99
interpreter_requires_environment)
1010
from test import support
11-
from test.support import force_not_colorized, warnings_helper
11+
from test.support import (force_not_colorized,
12+
force_not_colorized_test_class, warnings_helper)
1213
from test.support import os_helper
1314
from test.support import threading_helper
1415

@@ -1000,6 +1001,149 @@ def test_pymem_alloc0(self):
10001001
assert_python_ok('-X', 'tracemalloc', '-c', code)
10011002

10021003

1004+
@force_not_colorized_test_class
1005+
class TestModuleCommandLine(unittest.TestCase):
1006+
script = textwrap.dedent("""
1007+
import sys
1008+
def alloc():
1009+
return [bytes(10_000) for _ in range(100)]
1010+
data = alloc()
1011+
print(f'name={__name__} argv={sys.argv[1:]}')
1012+
""")
1013+
1014+
def setUp(self):
1015+
self.tmpdir = self.enterContext(os_helper.temp_dir())
1016+
self.script_path = self.write_script('script.py', self.script)
1017+
1018+
def write_script(self, name, source):
1019+
path = os.path.join(self.tmpdir, name)
1020+
with open(path, 'w', encoding='utf-8') as fp:
1021+
fp.write(source)
1022+
return path
1023+
1024+
def run_cli(self, *args, **kwargs):
1025+
res = assert_python_ok('-m', 'tracemalloc', *args, **kwargs)
1026+
return res.out.decode()
1027+
1028+
def run_cli_failure(self, *args):
1029+
res = assert_python_failure('-m', 'tracemalloc', *args)
1030+
return res.err.decode()
1031+
1032+
def test_script(self):
1033+
out = self.run_cli(self.script_path, 'a', '-b', '--c')
1034+
self.assertIn("name=__main__ argv=['a', '-b', '--c']", out)
1035+
self.assertIn('Top 10 lineno allocations still alive at exit:', out)
1036+
self.assertRegex(out, r'script\.py:4: size=\d+ KiB, count=\d+')
1037+
self.assertIn('Total size: ', out)
1038+
self.assertIn('Peak size: ', out)
1039+
1040+
def test_globals_alive(self):
1041+
# Memory referenced by globals is alive at exit, even without a
1042+
# reference cycle between a function and the globals
1043+
script = self.write_script(
1044+
'nofunc.py', 'data = [bytes(10_000) for _ in range(100)]\n')
1045+
out = self.run_cli(script)
1046+
self.assertRegex(out, r'nofunc\.py:1: size=\d+ KiB, count=\d+')
1047+
1048+
def test_module(self):
1049+
out = self.run_cli('-m', 'script', 'arg', __cwd=self.tmpdir,
1050+
__isolated=False)
1051+
self.assertIn("name=__main__ argv=['arg']", out)
1052+
self.assertRegex(out, r'script\.py:4: size=\d+ KiB')
1053+
1054+
def test_script_sys_path(self):
1055+
self.write_script('helper.py', 'print("helper imported")\n')
1056+
script = self.write_script('importer.py', 'import helper\n')
1057+
out = self.run_cli(script, __isolated=False)
1058+
self.assertIn('helper imported', out)
1059+
1060+
def test_limit(self):
1061+
out = self.run_cli('-l', '1', self.script_path)
1062+
lines = out.splitlines()
1063+
index = lines.index('Top 1 lineno allocations still alive at exit:')
1064+
self.assertRegex(lines[index + 1], r'script\.py:4: ')
1065+
self.assertRegex(lines[index + 2], r'^\d+ other: size=')
1066+
1067+
def test_key_type_filename(self):
1068+
out = self.run_cli('-t', 'filename', self.script_path)
1069+
self.assertRegex(out, r'script\.py:0: size=\d+ KiB')
1070+
1071+
def test_key_type_traceback(self):
1072+
out = self.run_cli('-t', 'traceback', '-n', '5', '-l', '1',
1073+
self.script_path)
1074+
self.assertIn('Top 1 traceback allocations still alive at exit:', out)
1075+
self.assertIn(' data = alloc()', out)
1076+
self.assertIn(' return [bytes(10_000) for _ in range(100)]', out)
1077+
# Frames of the command line interface are not displayed
1078+
self.assertNotIn('runpy', out)
1079+
self.assertNotIn('tracemalloc.py', out)
1080+
1081+
def test_cumulative(self):
1082+
out = self.run_cli('--cumulative', '-n', '5', '-l', '2',
1083+
self.script_path)
1084+
self.assertRegex(out, r'script\.py:5: size=\d+ KiB')
1085+
self.assertRegex(out, r'script\.py:4: size=\d+ KiB')
1086+
self.assertNotIn('runpy', out)
1087+
self.assertNotIn(' other: ', out)
1088+
1089+
def test_outfile(self):
1090+
outfile = os.path.join(self.tmpdir, 'snapshot.pickle')
1091+
out = self.run_cli('-n', '3', '-o', outfile, self.script_path)
1092+
self.assertNotIn('Total size', out)
1093+
snapshot = tracemalloc.Snapshot.load(outfile)
1094+
self.assertEqual(snapshot.traceback_limit, 3)
1095+
top = snapshot.statistics('lineno')[0]
1096+
self.assertEqual(top.traceback[0].filename, self.script_path)
1097+
self.assertEqual(top.traceback[0].lineno, 4)
1098+
1099+
def test_already_tracing(self):
1100+
# -n overrides the number of frames set by -X tracemalloc
1101+
outfile = os.path.join(self.tmpdir, 'snapshot.pickle')
1102+
self.run_cli('-n', '3', '-o', outfile, self.script_path,
1103+
__isolated=False, PYTHONTRACEMALLOC='1')
1104+
snapshot = tracemalloc.Snapshot.load(outfile)
1105+
self.assertEqual(snapshot.traceback_limit, 3)
1106+
1107+
def test_exit_code(self):
1108+
script = self.write_script('exit.py', 'import sys\nsys.exit(3)\n')
1109+
res = assert_python_failure('-m', 'tracemalloc', script)
1110+
self.assertEqual(res.rc, 3)
1111+
self.assertIn(b'Total size: ', res.out)
1112+
1113+
def test_exception(self):
1114+
script = self.write_script('error.py', '1/0\n')
1115+
res = assert_python_failure('-m', 'tracemalloc', script)
1116+
self.assertIn(b'ZeroDivisionError', res.err)
1117+
self.assertIn(b'Total size: ', res.out)
1118+
1119+
def test_tracing_stopped(self):
1120+
# The traced program can stop tracing
1121+
script = self.write_script(
1122+
'stop.py', 'import sys, tracemalloc\n'
1123+
'tracemalloc.stop()\n'
1124+
'sys.exit(3)\n')
1125+
res = assert_python_failure('-m', 'tracemalloc', script)
1126+
self.assertEqual(res.rc, 3)
1127+
self.assertEqual(res.out, b'')
1128+
self.assertEqual(res.err, b'')
1129+
1130+
def test_errors(self):
1131+
err = self.run_cli_failure()
1132+
self.assertIn('the following arguments are required: target', err)
1133+
err = self.run_cli_failure(os.path.join(self.tmpdir, 'missing.py'))
1134+
self.assertIn('cannot find script', err)
1135+
err = self.run_cli_failure('-t', 'traceback', '--cumulative',
1136+
self.script_path)
1137+
self.assertIn("--cumulative cannot be used with key type", err)
1138+
err = self.run_cli_failure('-l', '0', self.script_path)
1139+
self.assertIn('--limit must be a positive integer', err)
1140+
for nframe in INVALID_NFRAME:
1141+
with self.subTest(nframe=nframe):
1142+
err = self.run_cli_failure('-n', str(nframe),
1143+
self.script_path)
1144+
self.assertIn('the number of frames must be in range', err)
1145+
1146+
10031147
@unittest.skipIf(_testcapi is None, 'need _testcapi')
10041148
class TestCAPI(unittest.TestCase):
10051149
maxDiff = 80 * 20

‎Lib/tracemalloc.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -558,3 +558,138 @@ def take_snapshot():
558558
traces = _get_traces()
559559
traceback_limit = get_traceback_limit()
560560
return Snapshot(traces, traceback_limit)
561+
562+
563+
def _main(args=None):
564+
import argparse
565+
# Import pkgutil now, rather than when runpy imports it, so that its
566+
# memory is not attributed to the traced program.
567+
import pkgutil
568+
import runpy
569+
import sys
570+
571+
parser = argparse.ArgumentParser(
572+
prog='python -m tracemalloc',
573+
description='Trace memory allocations of a Python script or module '
574+
'and display the top allocations which are still alive '
575+
'when it finishes.',
576+
)
577+
parser.add_argument('-o', '--outfile',
578+
help='dump the snapshot to OUTFILE instead of '
579+
'displaying statistics; the file can be loaded '
580+
'with Snapshot.load()')
581+
parser.add_argument('-t', '--key-type', default='lineno',
582+
choices=('filename', 'lineno', 'traceback'),
583+
help='group statistics by key type '
584+
'(default: %(default)s)')
585+
parser.add_argument('--cumulative', action='store_true',
586+
help='cumulate the size and count of memory blocks '
587+
'of all frames of a traceback, not only the '
588+
'most recent frame')
589+
parser.add_argument('-n', '--nframe', type=int, default=1,
590+
help='maximum number of frames stored in a traceback '
591+
'(default: %(default)s)')
592+
parser.add_argument('-l', '--limit', type=int, default=10,
593+
help='number of statistics to display '
594+
'(default: %(default)s)')
595+
parser.add_argument('-m', dest='module', action='store_true',
596+
help='run target as a module (like python -m)')
597+
parser.add_argument('target',
598+
help='script file or module name to trace')
599+
parser.add_argument('args', nargs=argparse.REMAINDER,
600+
help='arguments passed to the script or module')
601+
opts = parser.parse_args(args)
602+
603+
if opts.cumulative and opts.key_type == 'traceback':
604+
parser.error("--cumulative cannot be used with key type 'traceback'")
605+
if opts.limit < 1:
606+
parser.error('--limit must be a positive integer')
607+
608+
# The traced program may chdir, so capture the absolute path now.
609+
if opts.outfile is not None:
610+
opts.outfile = os.path.abspath(opts.outfile)
611+
612+
if not opts.module:
613+
if not os.path.exists(opts.target):
614+
parser.error(f'cannot find script {opts.target!r}')
615+
# runpy.run_path() inserts directories and zip files in sys.path
616+
if (not sys.flags.safe_path
617+
and pkgutil.get_importer(opts.target) is None):
618+
sys.path[0] = os.path.dirname(os.path.realpath(opts.target))
619+
sys.argv[:] = [opts.target, *opts.args]
620+
621+
# Restart tracing if it was enabled by -X tracemalloc or
622+
# PYTHONTRACEMALLOC, so that the requested number of frames is used and
623+
# the report only covers the traced program.
624+
stop()
625+
try:
626+
start(opts.nframe)
627+
except ValueError as exc:
628+
parser.error(str(exc))
629+
630+
try:
631+
# Keep the globals of the traced program alive, so that the memory
632+
# they reference is included in the snapshot.
633+
if opts.module:
634+
globs = runpy.run_module(opts.target, run_name='__main__',
635+
alter_sys=True)
636+
else:
637+
globs = runpy.run_path(opts.target, run_name='__main__')
638+
finally:
639+
# The traced program can stop tracing
640+
if is_tracing():
641+
size, peak = get_traced_memory()
642+
snapshot = _strip_frames(take_snapshot(), {
643+
_main.__code__.co_filename,
644+
runpy._run_code.__code__.co_filename,
645+
})
646+
stop()
647+
if opts.outfile is not None:
648+
snapshot.dump(opts.outfile)
649+
else:
650+
_print_top(snapshot, opts.key_type, opts.cumulative,
651+
opts.limit, size, peak)
652+
653+
654+
def _strip_frames(snapshot, filenames):
655+
# Remove the oldest frames belonging to the given files (the frames of
656+
# the command line interface), but keep at least the most recent frame.
657+
if snapshot.traceback_limit == 1:
658+
return snapshot
659+
traces = []
660+
for domain, size, frames, total_nframe in snapshot.traces._traces:
661+
nframe = len(frames)
662+
while nframe > 1 and frames[nframe - 1][0] in filenames:
663+
nframe -= 1
664+
stripped = len(frames) - nframe
665+
traces.append((domain, size, frames[:nframe], total_nframe - stripped))
666+
return Snapshot(traces, snapshot.traceback_limit)
667+
668+
669+
def _print_top(snapshot, key_type, cumulative, limit, size, peak):
670+
stats = snapshot.statistics(key_type, cumulative)
671+
print(f'Top {min(limit, len(stats))} {key_type} allocations '
672+
f'still alive at exit:')
673+
for stat in stats[:limit]:
674+
if key_type == 'traceback':
675+
print(f'{stat.count} memory blocks: '
676+
f'{_format_size(stat.size, False)}')
677+
for line in stat.traceback.format():
678+
print(line)
679+
else:
680+
print(stat)
681+
if not cumulative and len(stats) > limit:
682+
other = stats[limit:]
683+
other_size = sum(stat.size for stat in other)
684+
other_count = sum(stat.count for stat in other)
685+
print(f'{len(other)} other: size={_format_size(other_size, False)}, '
686+
f'count={other_count}')
687+
print(f'Total size: {_format_size(size, False)}')
688+
print(f'Peak size: {_format_size(peak, False)}')
689+
690+
691+
if __name__ == '__main__':
692+
# Use the classes of the tracemalloc module rather than the ones of
693+
# __main__, so that dumped snapshots can be loaded.
694+
import tracemalloc
695+
tracemalloc._main()

0 commit comments

Comments
 (0)