Json formatter for python logger

json_formatter.py

import logging
import traceback

try:
    import ujson as json
except ImportError:
    import json


class JsonFormatter(logging.Formatter):

    TRACEBACK_TAGS = {
        0: 'file',
        1: 'line',
        2: 'fun',
        3: 'in'
    }

    def __init__(self, fields, datefmt=None, style=None):
        self._fields = fields
        self.datefmt = datefmt

    def usesTime(self):
        return 'asctime' in self._fields

    def formatMessage(self, record):
        record_dict = {
            field: getattr(record, field)
            for field in self._fields
            if hasattr(record, field)
        }
        return record_dict

    def formatStack(self, stack_info):
        return {'stack_info': str(stack_info)}

    def formatException(self, exc_info):
        exc_type, exc_value, exc_traceback = exc_info

        if isinstance(exc_traceback, traceback.StackSummary):
            tb = exc_traceback
        else:
            tb = traceback.extract_tb(exc_traceback)

        tb_tagged = [
            {self.TRACEBACK_TAGS[i]: f for i, f in enumerate(t)}
            for t in tb
        ]
        return {
            'exception': {
                'type': exc_type.__name__,
                'value': str(exc_value),
                'traceback': tb_tagged
            }
        }

    def format(self, record):
        if isinstance(record.msg, str):
            record.message = record.getMessage()
        else:
            record.message = record.msg

        if self.usesTime():
            record.asctime = self.formatTime(record, self.datefmt)

        s = self.formatMessage(record)

        if record.exc_info:
            if not hasattr(record, 'exc_dict'):
                record.exc_dict = self.formatException(record.exc_info)
        if hasattr(record, 'exc_dict'):
            s.update(record.exc_dict)
        if record.stack_info:
            s.update(self.formatStack(record.stack_info))

        return json.dumps(s)

test_json_formatter.py

import logging
import io
import json

import pytest
from pampy import match

from json_formatter import JsonFormatter


@pytest.fixture
def json_loggger():
    stream = io.StringIO()
    json_handler = logging.StreamHandler(stream)
    json_handler.setFormatter(JsonFormatter([
        'asctime',
        'levelname',
        'process',
        'name',
        'message',
        'context',
    ]))
    json_handler.setLevel(logging.INFO)

    logger = logging.getLogger()
    for handler in logger.handlers:
        logger.removeHandler(handler)

    logger.addHandler(json_handler)

    yield logger, stream

    json_handler.close()


def test_basic(json_loggger):
    logger, stream = json_loggger

    logger.info("Hello %s", "world")
    log_record = stream.getvalue()
    assert match(json.loads(log_record), {
        "asctime": str,
        "levelname": "INFO",
        "name": "root",
        "message": "Hello world"
    }, True, strict=False), log_record


def test_dict(json_loggger):
    logger, stream = json_loggger

    logger.info({"key": "value"})
    log_record = stream.getvalue()
    assert match(json.loads(log_record), {
        "asctime": str,
        "levelname": "INFO",
        "name": "root",
        "message": {"key": "value"}
    }, True, strict=False), log_record


def test_extra(json_loggger):
    logger, stream = json_loggger

    logger.info("Hello world", extra={"context": {"key": "value"}})
    log_record = stream.getvalue()
    assert match(json.loads(log_record), {
        "asctime": str,
        "levelname": "INFO",
        "name": "root",
        "message": "Hello world",
        "context": {"key": "value"}
    }, True, strict=False), log_record


def test_exception(json_loggger):
    logger, stream = json_loggger

    try:
        raise ValueError('incorrect value')
    except Exception:
        logger.exception('something wrong')

    log_record = stream.getvalue()
    assert match(json.loads(log_record), {
        "asctime": str,
        "levelname": "ERROR",
        "name": "root",
        "message": "something wrong",
        "exception": {
            "type": "ValueError",
            "value": "incorrect value",
            "traceback": [
                {
                    "file": __file__,
                    "line": int,
                    "fun": "test_exception",
                    "in": "raise ValueError('incorrect value')"
                }
            ]
        }
    }, True, strict=False), log_record