Add LogFormat enum

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -16,18 +16,15 @@
 from .lambda_runtime_log_utils import (
     JsonFormatter,
     _DATETIME_FORMAT,
-    _JSON_FORMAT,
+    LogFormat,
     _FRAME_TYPES,
     _DEFAULT_FRAME_TYPE,
-    _get_log_format_from_str,
 )
 from .lambda_runtime_marshaller import to_json
 
 ERROR_LOG_LINE_TERMINATE = "\r"
 ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
-_AWS_LAMBDA_LOG_FORMAT = _get_log_format_from_str(
-    os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT").upper()
-)
+_AWS_LAMBDA_LOG_FORMAT = LogFormat.from_str(os.environ.get("AWS_LAMBDA_LOG_FORMAT"))
 _AWS_LAMBDA_LOG_LEVEL = os.environ.get("AWS_LAMBDA_LOG_LEVEL")
 
 
@@ -109,7 +106,7 @@ def replace_line_indentation(line, indent_char, new_indent_char):
     return (new_indent_char * ident_chars_count) + line[ident_chars_count:]
 
 
-if _AWS_LAMBDA_LOG_FORMAT == _JSON_FORMAT:
+if _AWS_LAMBDA_LOG_FORMAT == LogFormat.JSON:
 
     def log_error(error_result, log_sink):
         error_result = {
@@ -423,7 +420,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_FORMAT:
+    if log_format == LogFormat.JSON:
         logger_handler.setFormatter(JsonFormatter())
     else:
         logger_handler.setFormatter(

awslambdaric/lambda_runtime_log_utils.py

@@ -5,6 +5,7 @@
 import logging
 import traceback
 import json
+from enum import IntEnum
 
 _DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
 _RESERVED_FIELDS = {
@@ -31,34 +32,39 @@
     "aws_request_id",
     "_frame_type",
 }
-_JSON_FORMAT = 0b0
-_TEXT_FORMAT = 0b1
+
+
+class LogFormat(IntEnum):
+    JSON = 0b0
+    TEXT = 0b1
+
+    @classmethod
+    def from_str(cls, value: str):
+        if value and value.upper() == "JSON":
+            return cls.JSON.value
+        return cls.TEXT.value
+
+
 _FRAME_TYPES = {
-    (_JSON_FORMAT, logging.NOTSET): 0xA55A0002.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.DEBUG): 0xA55A000A.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.INFO): 0xA55A000E.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.WARNING): 0xA55A0012.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.ERROR): 0xA55A0016.to_bytes(4, "big"),
-    (_JSON_FORMAT, logging.CRITICAL): 0xA55A001A.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.NOTSET): 0xA55A0003.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.DEBUG): 0xA55A000B.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.INFO): 0xA55A000F.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.WARNING): 0xA55A0013.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.ERROR): 0xA55A0017.to_bytes(4, "big"),
-    (_TEXT_FORMAT, logging.CRITICAL): 0xA55A001B.to_bytes(4, "big"),
+    (LogFormat.JSON, logging.NOTSET): 0xA55A0002.to_bytes(4, "big"),
+    (LogFormat.JSON, logging.DEBUG): 0xA55A000A.to_bytes(4, "big"),
+    (LogFormat.JSON, logging.INFO): 0xA55A000E.to_bytes(4, "big"),
+    (LogFormat.JSON, logging.WARNING): 0xA55A0012.to_bytes(4, "big"),
+    (LogFormat.JSON, logging.ERROR): 0xA55A0016.to_bytes(4, "big"),
+    (LogFormat.JSON, logging.CRITICAL): 0xA55A001A.to_bytes(4, "big"),
+    (LogFormat.TEXT, logging.NOTSET): 0xA55A0003.to_bytes(4, "big"),
+    (LogFormat.TEXT, logging.DEBUG): 0xA55A000B.to_bytes(4, "big"),
+    (LogFormat.TEXT, logging.INFO): 0xA55A000F.to_bytes(4, "big"),
+    (LogFormat.TEXT, logging.WARNING): 0xA55A0013.to_bytes(4, "big"),
+    (LogFormat.TEXT, logging.ERROR): 0xA55A0017.to_bytes(4, "big"),
+    (LogFormat.TEXT, logging.CRITICAL): 0xA55A001B.to_bytes(4, "big"),
 }
-_DEFAULT_FRAME_TYPE = _FRAME_TYPES[(_TEXT_FORMAT, logging.NOTSET)]
+_DEFAULT_FRAME_TYPE = _FRAME_TYPES[(LogFormat.TEXT, logging.NOTSET)]
 
 
 _encode_json = json.JSONEncoder(ensure_ascii=False).encode
 
 
-def _get_log_format_from_str(value: str):
-    if value == "JSON":
-        return _JSON_FORMAT
-    return _TEXT_FORMAT
-
-
 class JsonFormatter(logging.Formatter):
     def __init__(self):
         super().__init__(datefmt=_DATETIME_FORMAT)
@@ -98,8 +104,8 @@ def format_log_level(record: logging.LogRecord):
     def format(self, record: logging.LogRecord) -> str:
         self.format_log_level(record)
         record._frame_type = _FRAME_TYPES.get(
-            (_JSON_FORMAT, record.levelno),
-            _FRAME_TYPES[(_JSON_FORMAT, logging.NOTSET)],
+            (LogFormat.JSON, record.levelno),
+            _FRAME_TYPES[(LogFormat.JSON, logging.NOTSET)],
         )
 
         result = {

tests/test_bootstrap.py

@@ -19,11 +19,7 @@
 import awslambdaric.bootstrap as bootstrap
 from awslambdaric.lambda_runtime_exception import FaultException
 from awslambdaric.lambda_runtime_marshaller import LambdaMarshaller
-from awslambdaric.lambda_runtime_log_utils import (
-    _JSON_FORMAT,
-    _TEXT_FORMAT,
-    _get_log_format_from_str,
-)
+from awslambdaric.lambda_runtime_log_utils import LogFormat
 
 
 class TestUpdateXrayEnv(unittest.TestCase):
@@ -1178,7 +1174,7 @@ class TestLogging(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
         bootstrap.setup_logging(
-            _get_log_format_from_str("JSON"), "INFO", bootstrap.StandardLogSink()
+            LogFormat.from_str("JSON"), "INFO", bootstrap.StandardLogSink()
         )
 
     @patch("sys.stderr", new_callable=StringIO)