A Python decorator and custom logging handler to create call graphs

fun_logger.py

# 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')