Python traceback.walk_stack() Examples

def get_stack_info():
    '''Capture locals, module name, filename, and line number from the
    stacktrace to provide the source of the assertion error and
    formatted note.
    '''
    stack = traceback.walk_stack(sys._getframe().f_back)

    # We want locals from the test definition (which always begins
    # with 'test_' in unittest), which will be at a different
    # level in the stack depending on how many tests are in each
    # test case, how many test cases there are, etc.

    # The branch where we exhaust this loop is not covered
    # because we always find a test.
    for frame, _ in stack:  # pragma: no branch
        code = frame.f_code
        if code.co_name.startswith('test_'):
            return (frame.f_locals.copy(), frame.f_globals['__name__'],
                    code.co_filename, frame.f_lineno) 

def _log_trace(func, start=None):
    level = 0
    current_frame = None

    # Increase the current level for each traced function in the stackframe
    # This way we can visualize the call stack.
    for frame, _ in traceback.walk_stack(None):
        current_frame = current_frame if current_frame is not None else frame
        func_name = frame.f_code.co_name
        if func_name in _trace_funcs:
            level += 1

    # We can assume we have 'args' because we only call _log_trace inside
    # trace or TraceDispatcher.resgister
    current_op = current_frame.f_locals['args'][0]

    # If the first argument is a Expr, we print its op because it's more
    # informative.
    if isinstance(current_op, ir.Expr):
        current_op = current_op.op()

    _logger.debug(
        f"{'  ' * level} {func.__name__} {type(current_op).__qualname__} "
        f"{datetime.now() - start if start else ''}"
    ) 

def _extract_stack(limit=10):
    """Replacement for traceback.extract_stack() that only does the
    necessary work for asyncio debug mode.
    """
    frame = sys._getframe().f_back
    try:
        stack = traceback.StackSummary.extract(
            traceback.walk_stack(frame), lookup_lines=False)
    finally:
        del frame

    apg_path = asyncpg.__path__[0]
    i = 0
    while i < len(stack) and stack[i][0].startswith(apg_path):
        i += 1
    stack = stack[i:i + limit]

    stack.reverse()
    return ''.join(traceback.format_list(stack)) 

def __call__(self,*args,**kwargs):
        global RP_TRANSFERRING_FLAG
        if RP_TRANSFERRING_FLAG:
            return self.raw(*args,**kwargs)
        RP_TRANSFERRING_FLAG=True
        if not NET_INITTED:
            return self.raw(*args,**kwargs)
        for stack in traceback.walk_stack(None):
            if 'self' in stack[0].f_locals:
                layer=stack[0].f_locals['self']
                if layer in layer_names:
                    log.pytorch_layer_name=layer_names[layer]
                    print("Processing Layer: "+layer_names[layer])
                    break
        out=self.obj(self.raw,*args,**kwargs)
        RP_TRANSFERRING_FLAG=False
        # if isinstance(out,Variable):
        #     out=[out]
        return out

# ----- for torch.nn.functional operations ----- 

def test_extract_stack(self):
        s = traceback.StackSummary.extract(traceback.walk_stack(None))
        self.assertIsInstance(s, traceback.StackSummary) 

def extract_stack(f=None, limit=None):
    """Replacement for traceback.extract_stack() that only does the
    necessary work for asyncio debug mode.
    """
    if f is None:
        f = sys._getframe().f_back
    if limit is None:
        # Limit the amount of work to a reasonable amount, as extract_stack()
        # can be called for each coroutine and future in debug mode.
        limit = constants.DEBUG_STACK_DEPTH
    stack = traceback.StackSummary.extract(traceback.walk_stack(f),
                                           limit=limit,
                                           lookup_lines=False)
    stack.reverse()
    return stack 

def extract_stack(f=None, limit=None):
    """Replacement for traceback.extract_stack() that only does the
    necessary work for asyncio debug mode.
    """
    if f is None:
        f = sys._getframe().f_back
    if limit is None:
        # Limit the amount of work to a reasonable amount, as extract_stack()
        # can be called for each coroutine and future in debug mode.
        limit = constants.DEBUG_STACK_DEPTH
    stack = traceback.StackSummary.extract(traceback.walk_stack(f),
                                           limit=limit,
                                           lookup_lines=False)
    stack.reverse()
    return stack 

def _find_marker():
    marker_repr = repr(STACK_TRACE_MARKER)
    stack = traceback.StackSummary.extract(traceback.walk_stack(None), capture_locals=True)
    for frame in stack:
        if 'stack_trace_marker' in frame.locals:
            if frame.locals['stack_trace_marker'] == marker_repr:
                return True
    return False 

def _extract_stack_iter(frame):
        for f, lineno in walk_stack(frame):
            co = f.f_code
            filename = co.co_filename
            name = co.co_name
            yield filename, lineno, name 

def extract_stack(f=None, limit=None):
    """Replacement for traceback.extract_stack() that only does the
    necessary work for asyncio debug mode.
    """
    if f is None:
        f = sys._getframe().f_back
    if limit is None:
        # Limit the amount of work to a reasonable amount, as extract_stack()
        # can be called for each coroutine and future in debug mode.
        limit = constants.DEBUG_STACK_DEPTH
    stack = traceback.StackSummary.extract(traceback.walk_stack(f),
                                           limit=limit,
                                           lookup_lines=False)
    stack.reverse()
    return stack 

def test_format_locals(self):
        def some_inner(k, v):
            a = 1
            b = 2
            return traceback.StackSummary.extract(
                traceback.walk_stack(None), capture_locals=True, limit=1)
        s = some_inner(3, 4)
        self.assertEqual(
            ['  File "%s", line %d, in some_inner\n'
             '    traceback.walk_stack(None), capture_locals=True, limit=1)\n'
             '    a = 1\n'
             '    b = 2\n'
             '    k = 3\n'
             '    v = 4\n' % (__file__, some_inner.__code__.co_firstlineno + 4)
            ], s.format()) 

def test_extract_stack_limit(self):
        s = traceback.StackSummary.extract(traceback.walk_stack(None), limit=5)
        self.assertEqual(len(s), 5) 

def test_walk_stack(self):
        def deeper():
            return list(traceback.walk_stack(None))
        s1 = list(traceback.walk_stack(None))
        s2 = deeper()
        self.assertEqual(len(s2) - len(s1), 1)
        self.assertEqual(s2[1:], s1) 

def __call__(self, *args, **kwargs):
        if not NET_INITTED:
            return self.raw(*args, **kwargs)
        for stack in traceback.walk_stack(None):
            if 'self' in stack[0].f_locals:
                layer = stack[0].f_locals['self']
                if layer in layer_names:
                    log.pytorch_layer_name = layer_names[layer]
                    print(layer_names[layer])
                    break
        out = self.obj(self.raw, *args, **kwargs)
        # if isinstance(out,Variable):
        #     out=[out]
        return out 

def extract_stack(f=None, limit=None):
    """Replacement for traceback.extract_stack() that only does the
    necessary work for asyncio debug mode.
    """
    if f is None:
        f = sys._getframe().f_back
    if limit is None:
        # Limit the amount of work to a reasonable amount, as extract_stack()
        # can be called for each coroutine and future in debug mode.
        limit = constants.DEBUG_STACK_DEPTH
    stack = traceback.StackSummary.extract(traceback.walk_stack(f),
                                           limit=limit,
                                           lookup_lines=False)
    stack.reverse()
    return stack 

def __call__(self,*args,**kwargs):
        if not NET_INITTED:
            return self.raw(*args,**kwargs)
        for stack in traceback.walk_stack(None):
            if 'self' in stack[0].f_locals:
                layer=stack[0].f_locals['self']
                if layer in layer_names:
                    log.pytorch_layer_name=layer_names[layer]
                    print(layer_names[layer])
                    break
        out=self.obj(self.raw,*args,**kwargs)
        # if isinstance(out,Variable):
        #     out=[out]
        return out 

def test_format_locals(self):
        def some_inner(k, v):
            a = 1
            b = 2
            return traceback.StackSummary.extract(
                traceback.walk_stack(None), capture_locals=True, limit=1)
        s = some_inner(3, 4)
        self.assertEqual(
            ['  File "%s", line %d, in some_inner\n'
             '    traceback.walk_stack(None), capture_locals=True, limit=1)\n'
             '    a = 1\n'
             '    b = 2\n'
             '    k = 3\n'
             '    v = 4\n' % (__file__, some_inner.__code__.co_firstlineno + 4)
            ], s.format()) 

def test_extract_stack_limit(self):
        s = traceback.StackSummary.extract(traceback.walk_stack(None), limit=5)
        self.assertEqual(len(s), 5) 

def test_extract_stack(self):
        s = traceback.StackSummary.extract(traceback.walk_stack(None))
        self.assertIsInstance(s, traceback.StackSummary) 

def test_walk_stack(self):
        s = list(traceback.walk_stack(None))
        self.assertGreater(len(s), 10)