Refactor LambdaLogFormat

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -17,6 +17,7 @@
     INVOCATION_LOGGING_CONTEXT,
     JsonFormatter,
     DATETIME_FORMAT,
+    LambdaLogFormat,
 )
 from .lambda_runtime_marshaller import to_json
 
@@ -106,7 +107,7 @@ def replace_line_indentation(line, indent_char, new_indent_char):
 
 
 def log_error(error_result, log_sink):
-    if AWS_LAMBDA_LOG_FORMAT == "JSON":
+    if AWS_LAMBDA_LOG_FORMAT == LambdaLogFormat.JSON:
         log_error_json(error_result, log_sink)
     else:
         log_error_text(error_result, log_sink)
@@ -294,7 +295,7 @@ def emit(self, record: logging.LogRecord):
         self.log_sink.log(
             msg,
             log_level=record.levelno,
-            log_format=getattr(record, "log_format", "TEXT"),
+            log_format=getattr(record, "log_format", LambdaLogFormat.TEXT),
         )
 
 
@@ -336,7 +337,7 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         pass
 
-    def log(self, msg, log_level=logging.NOTSET, log_format="TEXT"):
+    def log(self, msg, log_level=logging.NOTSET, log_format=LambdaLogFormat.TEXT):
         sys.stdout.write(msg)
 
     def log_error(self, message_lines, log_level=logging.ERROR):
@@ -383,21 +384,11 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
     def set_log_level(self, frame_type, log_level):
         mask = self.LEVEL_TO_MASK.get(log_level, self.DEFAULT_LEVEL_MASK)
-        frame_type |= mask
-
-        return frame_type
-
-    @staticmethod
-    def set_log_format(frame_type, log_format):
-        if log_format == "JSON":
-            mask = 0b00
-        else:
-            mask = 0b01
         return frame_type | mask
 
-    def log(self, msg, log_level=logging.NOTSET, log_format="TEXT"):
+    def log(self, msg, log_level=logging.NOTSET, log_format=LambdaLogFormat.TEXT):
         frame_type = self.set_log_level(self.frame_type, log_level)
-        frame_type = self.set_log_format(frame_type, log_format)
+        frame_type = frame_type | log_format
 
         encoded_msg = msg.encode("utf8")
         timestamp = int(time.time_ns() / 1000)  # UNIX timestamp in microseconds
@@ -432,7 +423,9 @@ def create_log_sink():
         return StandardLogSink()
 
 
-AWS_LAMBDA_LOG_FORMAT = os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT")
+AWS_LAMBDA_LOG_FORMAT = LambdaLogFormat.from_str(
+    os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT")
+)
 AWS_LAMBDA_LOG_LEVEL = os.environ.get("AWS_LAMBDA_LOG_LEVEL")
 _GLOBAL_AWS_REQUEST_ID = None
 
@@ -442,7 +435,7 @@ def setup_logging(log_format, log_level, log_sink):
     logging.Formatter.converter = time.gmtime
     logger = logging.getLogger()
     logger_handler = LambdaLoggerHandler(log_sink)
-    if log_format == "JSON":
+    if log_format == LambdaLogFormat.JSON:
         logger_handler.setFormatter(JsonFormatter())
     else:
         logger_handler.setFormatter(

awslambdaric/lambda_runtime_log_utils.py

@@ -4,6 +4,7 @@
 
 import logging
 import traceback
+from enum import IntFlag
 
 from .lambda_runtime_marshaller import to_json
 
@@ -34,6 +35,17 @@
 }
 
 
+class LambdaLogFormat(IntFlag):
+    JSON = 0b0
+    TEXT = 0b1
+
+    @classmethod
+    def from_str(cls, value):
+        if value == "JSON":
+            return cls.JSON
+        return cls.TEXT
+
+
 class JsonFormatter(logging.Formatter):
     def __init__(self):
         super().__init__(datefmt=DATETIME_FORMAT, validate=False)
@@ -66,7 +78,7 @@ def format_location(record: logging.LogRecord):
         return f"{record.pathname}:{record.funcName}:{record.lineno}"
 
     def format(self, record: logging.LogRecord) -> str:
-        record.log_format = "JSON"
+        record.log_format = LambdaLogFormat.JSON
 
         result = {
             "timestamp": self.formatTime(record, self.datefmt),

tests/test_bootstrap.py

@@ -23,6 +23,7 @@
     lambda_invocation_logs_update,
     INVOCATION_LOGGING_CONTEXT,
     INSTANCE_LOGGING_CONTEXT,
+    LambdaLogFormat,
 )
 
 
@@ -1146,9 +1147,8 @@ def test_log_level_frame_type(self):
 
     def test_log_format_frame_type(self):
         test_cases = [
-            ("TEXT", 0xA55A0003),
-            ("WRONG_FORMAT", 0xA55A0003),
-            ("JSON", 0xA55A0002),
+            (LambdaLogFormat.TEXT, 0xA55A0003),
+            (LambdaLogFormat.JSON, 0xA55A0002),
         ]
 
         for fmt, expected_frame_type in test_cases:
@@ -1228,7 +1228,9 @@ def test_multiple_frame(self):
 class TestLogging(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
-        bootstrap.setup_logging("JSON", "INFO", bootstrap.StandardLogSink())
+        bootstrap.setup_logging(
+            LambdaLogFormat.from_str("JSON"), "INFO", bootstrap.StandardLogSink()
+        )
 
     def tearDown(self) -> None:
         INVOCATION_LOGGING_CONTEXT.clear()