From f11d8a4c7503add886509bd1dfafca5d3c2dd2e4 Mon Sep 17 00:00:00 2001
From: Daniel D'Avella <dan.davella@pixee.ai>
Date: Wed, 18 Oct 2023 13:37:52 -0400
Subject: [PATCH] Add optional project_name field to JSON log output

---
 src/codemodder/cli.py        |  4 ++++
 src/codemodder/codemodder.py |  2 +-
 src/codemodder/logging.py    | 16 ++++++++++++++--
 tests/test_logging.py        | 18 ++++++++++++++++++
 4 files changed, 37 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_logging.py

diff --git a/src/codemodder/cli.py b/src/codemodder/cli.py
index a48fa7deb..dddb68414 100644
--- a/src/codemodder/cli.py
+++ b/src/codemodder/cli.py
@@ -142,6 +142,10 @@ def parse_args(argv, codemod_registry):
         choices=[OutputFormat.HUMAN, OutputFormat.JSON],
         help="the format for the log output",
     )
+    parser.add_argument(
+        "--project-name",
+        help="optional descriptive name for the project used in log output",
+    )
     parser.add_argument(
         "--path-exclude",
         action=CsvListAction,
diff --git a/src/codemodder/codemodder.py b/src/codemodder/codemodder.py
index 6baef2fc7..a938e3876 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, argv.log_format)
+    configure_logger(argv.verbose, argv.log_format, argv.project_name)
 
     log_section("startup")
     logger.info("codemodder: python/%s", __VERSION__)
diff --git a/src/codemodder/logging.py b/src/codemodder/logging.py
index 5ea6dc18c..24769c688 100644
--- a/src/codemodder/logging.py
+++ b/src/codemodder/logging.py
@@ -1,6 +1,7 @@
 from enum import Enum
 import logging
 import sys
+from typing import Optional
 
 from pythonjsonlogger import jsonlogger
 
@@ -21,6 +22,12 @@ def __str__(self):
 
 
 class CodemodderJsonFormatter(jsonlogger.JsonFormatter):
+    project_name: Optional[str]
+
+    def __init__(self, *args, project_name: Optional[str] = None, **kwargs):
+        self.project_name = project_name
+        super().__init__(*args, **kwargs)
+
     def add_fields(self, log_record, record, message_dict):
         super().add_fields(log_record, record, message_dict)
         log_record["timestamp"] = log_record.pop("asctime")
@@ -28,6 +35,8 @@ def add_fields(self, log_record, record, message_dict):
         log_record["level"] = record.levelname.lower()
         log_record["file"] = record.filename
         log_record["line"] = record.lineno
+        if self.project_name:
+            log_record["project-name"] = self.project_name
 
 
 def log_section(section_name: str):
@@ -46,7 +55,9 @@ 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, log_format: OutputFormat):
+def configure_logger(
+    verbose: bool, log_format: OutputFormat, project_name: Optional[str] = None
+):
     """
     Configure the logger based on the verbosity level.
     """
@@ -64,7 +75,8 @@ def configure_logger(verbose: bool, log_format: OutputFormat):
             handlers.append(stderr_handler)
         case OutputFormat.JSON:
             formatter = CodemodderJsonFormatter(
-                "%(asctime) %(level) %(message) %(file) %(line)"
+                "%(asctime) %(level) %(message) %(file) %(line)",
+                project_name=project_name,
             )
             stdout_handler.setFormatter(formatter)
 
diff --git a/tests/test_logging.py b/tests/test_logging.py
new file mode 100644
index 000000000..caa56a73a
--- /dev/null
+++ b/tests/test_logging.py
@@ -0,0 +1,18 @@
+from pythonjsonlogger import jsonlogger
+
+from codemodder.logging import OutputFormat, configure_logger, logging
+
+
+def test_json_logger(mocker):
+    basic_config = mocker.patch("logging.basicConfig")
+    configure_logger(False, OutputFormat.JSON, "test-project")
+    assert basic_config.call_count == 1
+    assert basic_config.call_args[1]["format"] == "%(message)s"
+    assert isinstance(
+        basic_config.call_args[1]["handlers"][0].formatter,
+        jsonlogger.JsonFormatter,
+    )
+    assert (
+        basic_config.call_args[1]["handlers"][0].formatter.project_name
+        == "test-project"
+    )