Skip to content

Instantly share code, notes, and snippets.

@afvanwoudenberg
Created May 16, 2023 14:32
Show Gist options
  • Save afvanwoudenberg/0a460a8767e6a622519c8f3bf516e130 to your computer and use it in GitHub Desktop.
Save afvanwoudenberg/0a460a8767e6a622519c8f3bf516e130 to your computer and use it in GitHub Desktop.
A Python decorator and custom logging handler to create call graphs
# fun_logger.py
# For details see: https://www.aswinvanwoudenberg.com/posts/call-me-maybe/
import logging
import graphviz
from logging import StreamHandler
from functools import wraps, partial
def fun_logger(func=None, *, logging=logging, indent=' ', exit=False):
if func is None:
return partial(fun_logger,logging=logging, indent=indent, exit=exit)
global indent_level
indent_level = 0
@wraps(func)
def wrapper(*args, **kwargs):
global indent_level
args_repr = [repr(a) for a in args]
kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
signature = ", ".join(args_repr + kwargs_repr)
logging.debug(f"function {func.__qualname__} called with args {signature}", extra={
'func': func.__qualname__,
'args_': args,
'kwargs': kwargs,
'indent': indent_level * indent,
'indent_level': indent_level
})
try:
indent_level += 1
result = func(*args, **kwargs)
except Exception as e:
logging.exception(f"function {func.__qualname__} raised exception {str(e)}", extra={
'func': func.__qualname__,
'exception': e,
'indent': indent_level * indent,
'indent_level': indent_level
})
if indent_level > 0:
indent_level -= 1
raise e
if indent_level > 0:
indent_level -= 1
if exit:
logging.debug(f"function {func.__qualname__} exited with result {result}", extra={
'func': func.__qualname__,
'result': result,
'indent': indent_level * indent,
'indent_level': indent_level
})
return result
return wrapper
class CallGraphHandler(StreamHandler):
"""
A handler class which allows the drawing of call graphs
"""
def __init__(self, name=None, comment=None):
StreamHandler.__init__(self)
self.name = name
self.comment = comment
self.gv = graphviz.Digraph(name, comment)
self.clear()
@property
def source(self):
return self.gv.source
def clear(self):
self.gv.clear()
self.nodes = []
self.edges = []
self.stack = []
# insert empty node
self.gv.node('start', label='', shape='none')
def emit(self, record):
if hasattr(record, 'indent_level'):
current = None
while len(self.stack) > record.indent_level:
current = self.stack.pop()
if hasattr(record, 'args_'): # A function was called
args_repr = [repr(a) for a in record.args_]
kwargs_repr = [f"{k}={v!r}" for k, v in record.kwargs.items()]
signature = ", ".join(args_repr + kwargs_repr)
node_id = str(len(self.nodes))
self.gv.node(node_id, f"{record.func}({signature})", fontcolor='black', color='black')
self.nodes.append(node_id)
self.stack.append(node_id)
if len(self.stack) == 1:
self.gv.edge('start', node_id, fontcolor='black', color='black:black')
else:
self.gv.edge(self.stack[-2], node_id, fontcolor='black', color='black')
if hasattr(record, 'result'): # A function returned
if len(self.stack) == 0:
self.gv.edge(current, 'start', fontcolor='blue', pencolor='blue', labelfontcolor='blue', color='blue', label=str(record.result))
else:
self.gv.edge(current, self.stack[-1], fontcolor='blue', pencolor='blue', labelfontcolor='blue', color='blue', label=str(record.result))
if hasattr(record, 'exception'): # An exception occurred
if not current:
current = record.exception.__class__.__name__
self.gv.node(current, shape='diamond', fontcolor='red', pencolor='red', labelfontcolor='red', color='red')
self.gv.edge(self.stack[-1], current, fontcolor='red', pencolor='red', labelfontcolor='red', color='red')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment