import logging
from pythonjsonlogger import jsonlogger

class CustomJsonFormatter(jsonlogger.JsonFormatter):
    """As long as https://github.com/madzak/python-json-logger/issues/171 is still open, we need to have our
    own JsonFormatter implementation that fixes the bug."""

    def _perform_rename_log_fields(self, log_record):
        for old_field_name, new_field_name in self.rename_fields.items():
            if old_field_name not in log_record:
                continue
            log_record[new_field_name] = log_record[old_field_name]
            del log_record[old_field_name]

def _setup_logging():
    handler = logging.StreamHandler()
    handler.setFormatter(
        CustomJsonFormatter(
            "%(levelname)s %(message)s",
            rename_fields={"levelname": "severity", "exc_info": "stack_trace"},
        )
    )
    logging.getLogger().setLevel(logging.INFO)
    logging.getLogger().addHandler(handler)