From 04932c886c7f5683cc06ea22213593ed281aefb6 Mon Sep 17 00:00:00 2001
From: Daniel D'Avella <dan.davella@pixee.ai>
Date: Wed, 18 Oct 2023 11:49:08 -0400
Subject: [PATCH] Implement JSON logging with --log-format=json

---
 pyproject.toml               |  1 +
 src/codemodder/codemodder.py |  2 +-
 src/codemodder/logging.py    | 32 ++++++++++++++++++++++++++------
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 11dbd67c7..e130ae707 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,6 +13,7 @@ dependencies = [
     "isort~=5.12.0",
     "libcst~=1.1.0",
     "pylint~=3.0.0",
+    "python-json-logger~=2.0.0",
     "PyYAML~=6.0.0",
     "semgrep~=1.44.0",
     "wrapt~=1.15.0",
diff --git a/src/codemodder/codemodder.py b/src/codemodder/codemodder.py
index b25bcceeb..6baef2fc7 100644
--- a/src/codemodder/codemodder.py
+++ b/src/codemodder/codemodder.py
@@ -129,7 +129,7 @@ def run(original_args) -> int:
         )
         return 1
 
-    configure_logger(argv.verbose)
+    configure_logger(argv.verbose, argv.log_format)
 
     log_section("startup")
     logger.info("codemodder: python/%s", __VERSION__)
diff --git a/src/codemodder/logging.py b/src/codemodder/logging.py
index 5590b9dd5..5ea6dc18c 100644
--- a/src/codemodder/logging.py
+++ b/src/codemodder/logging.py
@@ -2,6 +2,8 @@
 import logging
 import sys
 
+from pythonjsonlogger import jsonlogger
+
 logger = logging.getLogger("codemodder")
 
 
@@ -18,6 +20,16 @@ def __str__(self):
         return self.value.lower()
 
 
+class CodemodderJsonFormatter(jsonlogger.JsonFormatter):
+    def add_fields(self, log_record, record, message_dict):
+        super().add_fields(log_record, record, message_dict)
+        log_record["timestamp"] = log_record.pop("asctime")
+        log_record.move_to_end("timestamp", last=False)
+        log_record["level"] = record.levelname.lower()
+        log_record["file"] = record.filename
+        log_record["line"] = record.lineno
+
+
 def log_section(section_name: str):
     """
     Log a section header.
@@ -34,22 +46,30 @@ def log_list(level: int, header: str, items: list, predicate=None):
         logger.log(level, "  - %s", predicate(item) if predicate else item)
 
 
-def configure_logger(verbose: bool):
+def configure_logger(verbose: bool, log_format: OutputFormat):
     """
     Configure the logger based on the verbosity level.
     """
     log_level = logging.DEBUG if verbose else logging.INFO
 
-    # TODO: this should all be conditional on the output format
     stdout_handler = logging.StreamHandler(sys.stdout)
     stdout_handler.setLevel(log_level)
-    stdout_handler.addFilter(lambda record: record.levelno <= logging.WARNING)
+    handlers = [stdout_handler]
 
-    stderr_handler = logging.StreamHandler(sys.stderr)
-    stderr_handler.setLevel(logging.ERROR)
+    match log_format:
+        case OutputFormat.HUMAN:
+            stdout_handler.addFilter(lambda record: record.levelno <= logging.WARNING)
+            stderr_handler = logging.StreamHandler(sys.stderr)
+            stderr_handler.setLevel(logging.ERROR)
+            handlers.append(stderr_handler)
+        case OutputFormat.JSON:
+            formatter = CodemodderJsonFormatter(
+                "%(asctime) %(level) %(message) %(file) %(line)"
+            )
+            stdout_handler.setFormatter(formatter)
 
     logging.basicConfig(
         format="%(message)s",
         level=log_level,
-        handlers=[stdout_handler, stderr_handler],
+        handlers=handlers,
     )