Address remaining suggestions of Andrew Svetlov

Author: ambv

Doc/library/asyncio-graph.rst

@@ -3,9 +3,9 @@
 
 .. _asyncio-graph:
 
-===================
-Stack Introspection
-===================
+========================
+Call Graph Introspection
+========================
 
 **Source code:** :source:`Lib/asyncio/graph.py`
 
@@ -20,20 +20,31 @@ and debuggers.
 .. versionadded:: next
 
 
-.. function:: print_call_graph(future=None, /, *, file=None, depth=1)
+.. function:: print_call_graph(future=None, /, *, file=None, depth=1, limit=None)
 
    Print the async call graph for the current task or the provided
    :class:`Task` or :class:`Future`.
 
+   This function prints entries starting from the currently executing frame,
+   i.e. the top frame, and going down towards the invocation point.
+
    The function receives an optional *future* argument.
-   If not passed, the current running task will be used. If there's no
-   current task, the function returns ``None``.
+   If not passed, the current running task will be used.
 
    If the function is called on *the current task*, the optional
    keyword-only *depth* argument can be used to skip the specified
    number of frames from top of the stack.
 
-   If *file* is omitted or ``None``, the function will print to :data:`sys.stdout`.
+   If the optional keyword-only *limit* argument is provided, each call stack
+   in the resulting graph is truncated to include at most ``abs(limit)``
+   entries. If *limit* is positive, the entries left are the closest to
+   the invocation point. If *limit* is negative, the topmost entries are
+   left. If *limit* is omitted or ``None``, all entries are present.
+   If *limit* is ``0``, the call stack is not printed at all, only
+   "awaited by" information is printed.
+
+   If *file* is omitted or ``None``, the function will print
+   to :data:`sys.stdout`.
 
    **Example:**
 
@@ -63,11 +74,14 @@ and debuggers.
             |   File 'taskgroups.py', line 107, in async TaskGroup.__aexit__()
             |   File 't2.py', line 7, in async main()
 
-.. function:: format_call_graph(future=None, /, *, depth=1)
+.. function:: format_call_graph(future=None, /, *, depth=1, limit=None)
 
    Like :func:`print_call_graph`, but returns a string.
+   If *future* is ``None`` and there's no current task,
+   the function returns an empty string.
+
 
-.. function:: capture_call_graph(future=None, /, *, depth=1)
+.. function:: capture_call_graph(future=None, /, *, depth=1, limit=None)
 
    Capture the async call graph for the current task or the provided
    :class:`Task` or :class:`Future`.

Lib/asyncio/graph.py

@@ -37,7 +37,11 @@ class FutureCallGraph:
     awaited_by: tuple["FutureCallGraph", ...]
 
 
-def _build_graph_for_future(future: futures.Future) -> FutureCallGraph:
+def _build_graph_for_future(
+    future: futures.Future,
+    *,
+    limit: int | None = None,
+) -> FutureCallGraph:
     if not isinstance(future, futures.Future):
         raise TypeError(
             f"{future!r} object does not appear to be compatible "
@@ -46,7 +50,7 @@ def _build_graph_for_future(future: futures.Future) -> FutureCallGraph:
 
     coro = None
     if get_coro := getattr(future, 'get_coro', None):
-        coro = get_coro()
+        coro = get_coro() if limit != 0 else None
 
     st: list[FrameCallGraphEntry] = []
     awaited_by: list[FutureCallGraph] = []
@@ -65,8 +69,13 @@ def _build_graph_for_future(future: futures.Future) -> FutureCallGraph:
 
     if future._asyncio_awaited_by:
         for parent in future._asyncio_awaited_by:
-            awaited_by.append(_build_graph_for_future(parent))
+            awaited_by.append(_build_graph_for_future(parent, limit=limit))
 
+    if limit is not None:
+        if limit > 0:
+            st = st[:limit]
+        elif limit < 0:
+            st = st[limit:]
     st.reverse()
     return FutureCallGraph(future, tuple(st), tuple(awaited_by))
 
@@ -76,8 +85,9 @@ def capture_call_graph(
     /,
     *,
     depth: int = 1,
+    limit: int | None = None,
 ) -> FutureCallGraph | None:
-    """Capture async call graph for the current task or the provided Future.
+    """Capture the async call graph for the current task or the provided Future.
 
     The graph is represented with three data structures:
 
@@ -95,13 +105,21 @@ def capture_call_graph(
       Where 'frame' is a frame object of a regular Python function
       in the call stack.
 
-    Receives an optional "future" argument. If not passed,
+    Receives an optional 'future' argument. If not passed,
     the current task will be used. If there's no current task, the function
     returns None.
 
     If "capture_call_graph()" is introspecting *the current task*, the
-    optional keyword-only "depth" argument can be used to skip the specified
+    optional keyword-only 'depth' argument can be used to skip the specified
     number of frames from top of the stack.
+
+    If the optional keyword-only 'limit' argument is provided, each call stack
+    in the resulting graph is truncated to include at most ``abs(limit)``
+    entries. If 'limit' is positive, the entries left are the closest to
+    the invocation point. If 'limit' is negative, the topmost entries are
+    left. If 'limit' is omitted or None, all entries are present.
+    If 'limit' is 0, the call stack is not captured at all, only
+    "awaited by" information is present.
     """
 
     loop = events._get_running_loop()
@@ -111,7 +129,7 @@ def capture_call_graph(
         # if yes - check if the passed future is the currently
         # running task or not.
         if loop is None or future is not tasks.current_task(loop=loop):
-            return _build_graph_for_future(future)
+            return _build_graph_for_future(future, limit=limit)
         # else: future is the current task, move on.
     else:
         if loop is None:
@@ -134,7 +152,7 @@ def capture_call_graph(
 
     call_stack: list[FrameCallGraphEntry] = []
 
-    f = sys._getframe(depth)
+    f = sys._getframe(depth) if limit != 0 else None
     try:
         while f is not None:
             is_async = f.f_generator is not None
@@ -153,7 +171,14 @@ def capture_call_graph(
     awaited_by = []
     if future._asyncio_awaited_by:
         for parent in future._asyncio_awaited_by:
-            awaited_by.append(_build_graph_for_future(parent))
+            awaited_by.append(_build_graph_for_future(parent, limit=limit))
+
+    if limit is not None:
+        limit *= -1
+        if limit > 0:
+            call_stack = call_stack[:limit]
+        elif limit < 0:
+            call_stack = call_stack[limit:]
 
     return FutureCallGraph(future, tuple(call_stack), tuple(awaited_by))
 
@@ -163,8 +188,9 @@ def format_call_graph(
     /,
     *,
     depth: int = 1,
+    limit: int | None = None,
 ) -> str:
-    """Return async call graph as a string for `future`.
+    """Return the async call graph as a string for `future`.
 
     If `future` is not provided, format the call graph for the current task.
     """
@@ -226,9 +252,9 @@ def add_line(line: str) -> None:
             for fut in st.awaited_by:
                 render_level(fut, buf, level + 1)
 
-    graph = capture_call_graph(future, depth=depth + 1)
+    graph = capture_call_graph(future, depth=depth + 1, limit=limit)
     if graph is None:
-        return
+        return ""
 
     try:
         buf: list[str] = []
@@ -245,6 +271,7 @@ def print_call_graph(
     *,
     file: typing.TextIO | None = None,
     depth: int = 1,
+    limit: int | None = None,
 ) -> None:
-    """Print async call graph for the current task or the provided Future."""
-    print(format_call_graph(future, depth=depth), file=file)
+    """Print the async call graph for the current task or the provided Future."""
+    print(format_call_graph(future, depth=depth, limit=limit), file=file)