Issue #18874: Remove tracemalloc.set_traceback_limit()

tracemalloc.start() now has an option nframe parameter

4 files changed, 52 insertions, 72 deletions

diff --git a/Doc/library/tracemalloc.rst b/Doc/library/tracemalloc.rst
index 2edf7bb..0e11540 100644
--- a/Doc/library/tracemalloc.rst
+++ b/Doc/library/tracemalloc.rst
@@ -21,8 +21,7 @@ start tracing Python memory allocations.
 By default, a trace of an allocated memory block only stores the most recent
 frame (1 frame). To store 25 frames at startup: set the
 :envvar:`PYTHONTRACEMALLOC` environment variable to ``25``, or use the
-:option:`-X` ``tracemalloc=25`` command line option. The
-:func:`set_traceback_limit` function can be used at runtime to set the limit.
+:option:`-X` ``tracemalloc=25`` command line option.
 
 .. versionadded:: 3.4
 
@@ -120,8 +119,8 @@ Code to display the traceback of the biggest memory block::
     import linecache
     import tracemalloc
 
-    tracemalloc.set_traceback_limit(25)
-    tracemalloc.start()
+    # Store 25 frames
+    tracemalloc.start(25)
 
     # ... run your application ...
 
@@ -267,10 +266,10 @@ Functions
 
    Get the maximum number of frames stored in the traceback of a trace.
 
-   By default, a trace of a memory block only stores the most recent
-   frame: the limit is ``1``.
+   The :mod:`tracemalloc` module must be tracing memory allocations to
+   get the limit, otherwise an exception is raised.
 
-   Use the :func:`set_traceback_limit` function to change the limit.
+   The limit is set by the :func:`start` function.
 
 
 .. function:: get_traced_memory()
@@ -294,10 +293,12 @@ Functions
     See also :func:`start` and :func:`stop` functions.
 
 
-.. function:: set_traceback_limit(nframe: int)
+.. function:: start(nframe: int=1)
 
-   Set the maximum number of frames stored in the traceback of a trace.
-   *nframe* must be greater or equal to ``1``.
+   Start tracing Python memory allocations: install hooks on Python memory
+   allocators. Collected tracebacks of traces will be limited to *nframe*
+   frames. By default, a trace of a memory block only stores the most recent
+   frame: the limit is ``1``. *nframe* must be greater or equal to ``1``.
 
    Storing more than ``1`` frame is only useful to compute statistics grouped
    by ``'traceback'`` or to compute cumulative statistics: see the
@@ -309,17 +310,10 @@ Functions
 
    The :envvar:`PYTHONTRACEMALLOC` environment variable
    (``PYTHONTRACEMALLOC=NFRAME``) and the :option:`-X` ``tracemalloc=NFRAME``
-   command line option can be used to set the limit at startup.
-
-   Use the :func:`get_traceback_limit` function to get the current limit.
-
-
-.. function:: start()
-
-   Start tracing Python memory allocations: install hooks on Python memory
-   allocators.
+   command line option can be used to start tracing at startup.
 
-   See also :func:`stop` and :func:`is_tracing` functions.
+   See also :func:`stop`, :func:`is_tracing` and :func:`get_traceback_limit`
+   functions.
 
 
 .. function:: stop()
@@ -342,7 +336,7 @@ Functions
    :mod:`tracemalloc` module started to trace memory allocations.
 
    Tracebacks of traces are limited to :func:`get_traceback_limit` frames. Use
-   :func:`set_traceback_limit` to store more frames.
+   the *nframe* parameter of the :func:`start` function to store more frames.
 
    The :mod:`tracemalloc` module must be tracing memory allocations to take a
    snapshot, see the the :func:`start` function.
diff --git a/Doc/using/cmdline.rst b/Doc/using/cmdline.rst
index cfd6045..d1822fb 100644
--- a/Doc/using/cmdline.rst
+++ b/Doc/using/cmdline.rst
@@ -381,10 +381,11 @@ Miscellaneous options
    * ``-X faulthandler`` to enable :mod:`faulthandler`;
    * ``-X showrefcount`` to enable the output of the total reference count
      and memory blocks (only works on debug builds);
-   * ``-X tracemalloc`` to enable :mod:`tracemalloc`.
-   * ``-X tracemalloc=NFRAME`` to enable :mod:`tracemalloc`, *NFRAME* is the
-     maximum number of frames stored in a trace: see the
-     :func:`tracemalloc.set_traceback_limit` function.
+   * ``-X tracemalloc`` to start tracing Python memory allocations using the
+     :mod:`tracemalloc` module. By default, only the most recent frame is
+     stored in a traceback of a trace. Use ``-X tracemalloc=NFRAME`` to start
+     tracing with a traceback limit of *NFRAME* frames. See the
+     :func:`tracemalloc.start` for more information.
 
    It also allows to pass arbitrary values and retrieve them through the
    :data:`sys._xoptions` dictionary.
@@ -600,10 +601,11 @@ conflict.
 
 .. envvar:: PYTHONTRACEMALLOC
 
-   If this environment variable is set to a non-empty string, all memory
-   allocations made by Python are traced by the :mod:`tracemalloc` module.
-   The value of the variable is the maximum number of frames stored in a trace:
-   see the :func:`tracemalloc.set_traceback_limit` function.
+   If this environment variable is set to a non-empty string, start tracing
+   Python memory allocations using the :mod:`tracemalloc` module. The value of
+   the variable is the maximum number of frames stored in a traceback of a
+   trace. For example, ``PYTHONTRACEMALLOC=1`` stores only the most recent
+   frame. See the :func:`tracemalloc.start` for more information.
 
    .. versionadded:: 3.4
 
diff --git a/Lib/test/test_tracemalloc.py b/Lib/test/test_tracemalloc.py
index ecb5aee..1d15667 100644
--- a/Lib/test/test_tracemalloc.py
+++ b/Lib/test/test_tracemalloc.py
@@ -83,8 +83,7 @@ class TestTracemallocEnabled(unittest.TestCase):
         if tracemalloc.is_tracing():
             self.skipTest("tracemalloc must be stopped before the test")
 
-        tracemalloc.set_traceback_limit(1)
-        tracemalloc.start()
+        tracemalloc.start(1)
 
     def tearDown(self):
         tracemalloc.stop()
@@ -109,20 +108,18 @@ class TestTracemallocEnabled(unittest.TestCase):
     def test_set_traceback_limit(self):
         obj_size = 10
 
-        nframe = tracemalloc.get_traceback_limit()
-        self.addCleanup(tracemalloc.set_traceback_limit, nframe)
-
-        self.assertRaises(ValueError, tracemalloc.set_traceback_limit, -1)
+        tracemalloc.stop()
+        self.assertRaises(ValueError, tracemalloc.start, -1)
 
-        tracemalloc.clear_traces()
-        tracemalloc.set_traceback_limit(10)
+        tracemalloc.stop()
+        tracemalloc.start(10)
         obj2, obj2_traceback = allocate_bytes(obj_size)
         traceback = tracemalloc.get_object_traceback(obj2)
         self.assertEqual(len(traceback), 10)
         self.assertEqual(traceback, obj2_traceback)
 
-        tracemalloc.clear_traces()
-        tracemalloc.set_traceback_limit(1)
+        tracemalloc.stop()
+        tracemalloc.start(1)
         obj, obj_traceback = allocate_bytes(obj_size)
         traceback = tracemalloc.get_object_traceback(obj)
         self.assertEqual(len(traceback), 1)
@@ -163,8 +160,8 @@ class TestTracemallocEnabled(unittest.TestCase):
             return allocate_bytes3(size)
 
         # Ensure that two identical tracebacks are not duplicated
-        tracemalloc.clear_traces()
-        tracemalloc.set_traceback_limit(4)
+        tracemalloc.stop()
+        tracemalloc.start(4)
         obj_size = 123
         obj1, obj1_traceback = allocate_bytes4(obj_size)
         obj2, obj2_traceback = allocate_bytes4(obj_size)
diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c
index 15ed734..9dd4b19 100644
--- a/Modules/_tracemalloc.c
+++ b/Modules/_tracemalloc.c
@@ -1151,13 +1151,27 @@ done:
 }
 
 PyDoc_STRVAR(tracemalloc_start_doc,
-    "start()\n"
+    "start(nframe: int=1)\n"
     "\n"
-    "Start tracing Python memory allocations.");
+    "Start tracing Python memory allocations. Set also the maximum number \n"
+    "of frames stored in the traceback of a trace to nframe.");
 
 static PyObject*
-py_tracemalloc_start(PyObject *self)
+py_tracemalloc_start(PyObject *self, PyObject *args)
 {
+    Py_ssize_t nframe = 1;
+
+    if (!PyArg_ParseTuple(args, "|n:start", &nframe))
+        return NULL;
+
+    if (nframe < 1 || nframe > MAX_NFRAME) {
+        PyErr_Format(PyExc_ValueError,
+                     "the number of frames must be in range [1; %i]",
+                     MAX_NFRAME);
+        return NULL;
+    }
+    tracemalloc_config.max_nframe = Py_SAFE_DOWNCAST(nframe, Py_ssize_t, int);
+
     if (tracemalloc_start() < 0)
         return NULL;
 
@@ -1192,31 +1206,6 @@ py_tracemalloc_get_traceback_limit(PyObject *self)
     return PyLong_FromLong(tracemalloc_config.max_nframe);
 }
 
-PyDoc_STRVAR(tracemalloc_set_traceback_limit_doc,
-    "set_traceback_limit(nframe: int)\n"
-    "\n"
-    "Set the maximum number of frames stored in the traceback of a trace.");
-
-static PyObject*
-tracemalloc_set_traceback_limit(PyObject *self, PyObject *args)
-{
-    Py_ssize_t nframe;
-
-    if (!PyArg_ParseTuple(args, "n:set_traceback_limit",
-                          &nframe))
-        return NULL;
-
-    if (nframe < 1 || nframe > MAX_NFRAME) {
-        PyErr_Format(PyExc_ValueError,
-                     "the number of frames must be in range [1; %i]",
-                     MAX_NFRAME);
-        return NULL;
-    }
-    tracemalloc_config.max_nframe = Py_SAFE_DOWNCAST(nframe, Py_ssize_t, int);
-
-    Py_RETURN_NONE;
-}
-
 PyDoc_STRVAR(tracemalloc_get_tracemalloc_memory_doc,
     "get_tracemalloc_memory() -> int\n"
     "\n"
@@ -1275,13 +1264,11 @@ static PyMethodDef module_methods[] = {
     {"_get_object_traceback", (PyCFunction)py_tracemalloc_get_object_traceback,
      METH_O, tracemalloc_get_object_traceback_doc},
     {"start", (PyCFunction)py_tracemalloc_start,
-      METH_NOARGS, tracemalloc_start_doc},
+      METH_VARARGS, tracemalloc_start_doc},
     {"stop", (PyCFunction)py_tracemalloc_stop,
       METH_NOARGS, tracemalloc_stop_doc},
     {"get_traceback_limit", (PyCFunction)py_tracemalloc_get_traceback_limit,
      METH_NOARGS, tracemalloc_get_traceback_limit_doc},
-    {"set_traceback_limit", (PyCFunction)tracemalloc_set_traceback_limit,
-     METH_VARARGS, tracemalloc_set_traceback_limit_doc},
     {"get_tracemalloc_memory", (PyCFunction)tracemalloc_get_tracemalloc_memory,
      METH_NOARGS, tracemalloc_get_tracemalloc_memory_doc},
     {"get_traced_memory", (PyCFunction)tracemalloc_get_traced_memory,