Make stack_trace() shortcut use default program argument logic

Instead of always using the default program, if the thread argument is an Object, we should use thread.prog_. Signed-off-by: Omar Sandoval <[email protected]>
osandov · Dec 1, 2023 · 433c735 · 433c735
1 parent 8ed2403
commit 433c735
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/drgn/__init__.py b/drgn/__init__.py
@@ -44,6 +44,7 @@
 import pkgutil
 import sys
 import types
+from typing import Union
 
 from _drgn import (
     NULL,
@@ -263,11 +264,17 @@ def task_exe_path(task):
             del sys.modules["__main__"]
 
 
-def stack_trace(thread: Object) -> StackTrace:
+def stack_trace(thread: Union[Object, IntegerLike]) -> StackTrace:
     """
-    Get the stack trace for the given thread in the :ref:`default-program`.
+    Get the stack trace for the given thread using the :ref:`default program
+    argument <default-program>`.
 
-    This is equivalent to ``get_default_prog().stack_trace(thread)``. See
-    :meth:`Program.stack_trace()` for more details.
+    See :meth:`Program.stack_trace()` for more details.
+
+    :param thread: Thread ID, ``struct pt_regs`` object, or
+        ``struct task_struct *`` object.
     """
-    return get_default_prog().stack_trace(thread)
+    if isinstance(thread, Object):
+        return thread.prog_.stack_trace(thread)
+    else:
+        return get_default_prog().stack_trace(thread)