feat: Faster JSON dumps in format_message with msgspec

pyproject.toml

@@ -49,6 +49,7 @@ inflection = ">=0.5.1"
 joblib = ">=1.3.0"
 jsonpath-ng = ">=1.5.3"
 jsonschema = ">=4.16.0"
+msgspec = ">=0.18.0"
 packaging = ">=23.1"
 pendulum = ">=2.1.0,<4"
 python-dateutil = ">=2.8.2"

singer_sdk/helpers/_flattening.py

@@ -7,13 +7,29 @@
 import re
 import typing as t
 from copy import deepcopy
+from datetime import datetime
 
 import inflection
-import simplejson as json
+import msgspec
 
 DEFAULT_FLATTENING_SEPARATOR = "__"
 
 
+def enc_hook(obj: t.Any) -> t.Any:  # noqa: ANN401
+    """Enocding type helper for non native types.
+
+    Args:
+        obj: the item to be encoded
+
+    Returns:
+        The object converted to the appropriate type, default is str
+    """
+    return obj.isoformat(sep="T") if isinstance(obj, datetime) else str(obj)
+
+
+encoder = msgspec.json.Encoder(enc_hook=enc_hook, decimal_format="number")
+
+
 class FlatteningOptions(t.NamedTuple):
     """A stream map which performs the flattening role."""
 
@@ -435,7 +451,7 @@ def _flatten_record(
             items.append(
                 (
                     new_key,
-                    json.dumps(v, use_decimal=True, default=str)
+                    encoder.encode(v).decode()
                     if _should_jsondump_value(k, v, flattened_schema)
                     else v,
                 ),

singer_sdk/io_base.py

@@ -4,25 +4,57 @@
 
 import abc
 import decimal
-import json
 import logging
 import sys
 import typing as t
 from collections import Counter, defaultdict
+from datetime import datetime
+
+import msgspec
 
 from singer_sdk._singerlib.messages import Message, SingerMessageType
-from singer_sdk._singerlib.messages import format_message as singer_format_message
-from singer_sdk._singerlib.messages import write_message as singer_write_message
 from singer_sdk.exceptions import InvalidInputLine
 
 logger = logging.getLogger(__name__)
+msg_buffer = bytearray(64)
+
+
+def enc_hook(obj: t.Any) -> t.Any:  # noqa: ANN401
+    """Enocding type helper for non native types.
+
+    Args:
+        obj: the item to be encoded
+
+    Returns:
+        The object converted to the appropriate type, default is str
+    """
+    return obj.isoformat(sep="T") if isinstance(obj, datetime) else str(obj)
+
+
+encoder = msgspec.json.Encoder(enc_hook=enc_hook, decimal_format="number")
+
+
+def dec_hook(type: type, obj: t.Any) -> t.Any:  # noqa: ARG001, A002, ANN401
+    """Decoding type helper for non native types.
+
+    Args:
+        type: the type given
+        obj: the item to be decoded
+
+    Returns:
+        The object converted to the appropriate type, default is str.
+    """
+    return str(obj)
+
+
+decoder = msgspec.json.Decoder(dec_hook=dec_hook, float_hook=decimal.Decimal)
 
 
 class SingerReader(metaclass=abc.ABCMeta):
     """Interface for all plugins reading Singer messages from stdin."""
 
     @t.final
-    def listen(self, file_input: t.IO[str] | None = None) -> None:
+    def listen(self, file_input: t.IO | None = None) -> None:
         """Read from input until all messages are processed.
 
         Args:
@@ -31,7 +63,7 @@ def listen(self, file_input: t.IO[str] | None = None) -> None:
         This method is internal to the SDK and should not need to be overridden.
         """
         if not file_input:
-            file_input = sys.stdin
+            file_input = sys.stdin.buffer
 
         self._process_lines(file_input)
         self._process_endofpipe()
@@ -52,7 +84,7 @@ def _assert_line_requires(line_dict: dict, requires: set[str]) -> None:
             msg = f"Line is missing required {', '.join(missing)} key(s): {line_dict}"
             raise InvalidInputLine(msg)
 
-    def deserialize_json(self, line: str) -> dict:  # noqa: PLR6301
+    def deserialize_json(self, line: bytes) -> dict:  # noqa: PLR6301
         """Deserialize a line of json.
 
         Args:
@@ -62,18 +94,17 @@ def deserialize_json(self, line: str) -> dict:  # noqa: PLR6301
             A dictionary of the deserialized json.
 
         Raises:
-            json.decoder.JSONDecodeError: raised if any lines are not valid json
+            msgspec.DecodeError: raised if any lines are not valid json
         """
         try:
-            return json.loads(  # type: ignore[no-any-return]
+            return decoder.decode(  # type: ignore[no-any-return]
                 line,
-                parse_float=decimal.Decimal,
             )
-        except json.decoder.JSONDecodeError as exc:
+        except msgspec.DecodeError as exc:
             logger.exception("Unable to parse:\n%s", line, exc_info=exc)
             raise
 
-    def _process_lines(self, file_input: t.IO[str]) -> t.Counter[str]:
+    def _process_lines(self, file_input: t.IO) -> t.Counter[str]:
         """Internal method to process jsonl lines from a Singer tap.
 
         Args:
@@ -145,7 +176,7 @@ def _process_endofpipe(self) -> None:  # noqa: PLR6301
 class SingerWriter:
     """Interface for all plugins writting Singer messages to stdout."""
 
-    def format_message(self, message: Message) -> str:  # noqa: PLR6301
+    def format_message(self, message: Message) -> bytes:  # noqa: PLR6301
         """Format a message as a JSON string.
 
         Args:
@@ -154,12 +185,15 @@ def format_message(self, message: Message) -> str:  # noqa: PLR6301
         Returns:
             The formatted message.
         """
-        return singer_format_message(message)
+        encoder.encode_into(message.to_dict(), msg_buffer)
+        msg_buffer.extend(b"\n")
+        return msg_buffer
 
-    def write_message(self, message: Message) -> None:  # noqa: PLR6301
+    def write_message(self, message: Message) -> None:
         """Write a message to stdout.
 
         Args:
             message: The message to write.
         """
-        singer_write_message(message)
+        sys.stdout.buffer.write(self.format_message(message))
+        sys.stdout.flush()

singer_sdk/mapper_base.py

@@ -112,7 +112,7 @@ def invoke(  # type: ignore[override]
         about: bool = False,
         about_format: str | None = None,
         config: tuple[str, ...] = (),
-        file_input: t.IO[str] | None = None,
+        file_input: t.IO[bytes] | None = None,
     ) -> None:
         """Invoke the mapper.
 

singer_sdk/target_base.py

@@ -294,7 +294,7 @@ def _handle_max_record_age(self) -> None:
             )
             self.drain_all()
 
-    def _process_lines(self, file_input: t.IO[str]) -> t.Counter[str]:
+    def _process_lines(self, file_input: t.IO) -> t.Counter[str]:
         """Internal method to process jsonl lines from a Singer tap.
 
         Args:
@@ -544,7 +544,7 @@ def invoke(  # type: ignore[override]
         about: bool = False,
         about_format: str | None = None,
         config: tuple[str, ...] = (),
-        file_input: t.IO[str] | None = None,
+        file_input: t.IO[bytes] | None = None,
     ) -> None:
         """Invoke the target.
 

singer_sdk/testing/legacy.py

@@ -7,6 +7,7 @@
 from contextlib import redirect_stderr, redirect_stdout
 
 import singer_sdk._singerlib as singer
+from singer_sdk.testing.runners import RunnerStandardOutErr
 
 if t.TYPE_CHECKING:
     from singer_sdk.mapper_base import InlineMapper
@@ -111,7 +112,7 @@ def get_standard_target_tests(
     return []
 
 
-def tap_sync_test(tap: Tap) -> tuple[io.StringIO, io.StringIO]:
+def tap_sync_test(tap: Tap) -> tuple[RunnerStandardOutErr, io.StringIO]:
     """Invokes a Tap object and return STDOUT and STDERR results in StringIO buffers.
 
     Args:
@@ -120,11 +121,15 @@ def tap_sync_test(tap: Tap) -> tuple[io.StringIO, io.StringIO]:
     Returns:
         A 2-item tuple with StringIO buffers from the Tap's output: (stdout, stderr)
     """
-    stdout_buf = io.StringIO()
+    stdout_buf = RunnerStandardOutErr()
     stderr_buf = io.StringIO()
+
     with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf):
         tap.sync_all()
-    stdout_buf.seek(0)
+
+    # Add decoded buffer items into stdout_buf
+    stdout_buf.load_from_buffer()
+
     stderr_buf.seek(0)
     return stdout_buf, stderr_buf
 
@@ -171,7 +176,7 @@ def _select_all(catalog_dict: dict) -> dict:
 
 def target_sync_test(
     target: Target,
-    input: io.StringIO | None,  # noqa: A002
+    input: RunnerStandardOutErr | None,  # noqa: A002
     *,
     finalize: bool = True,
 ) -> tuple[io.StringIO, io.StringIO]:
@@ -203,7 +208,7 @@ def target_sync_test(
 def tap_to_target_sync_test(
     tap: Tap,
     target: Target,
-) -> tuple[io.StringIO, io.StringIO, io.StringIO, io.StringIO]:
+) -> tuple[RunnerStandardOutErr, io.StringIO, io.StringIO, io.StringIO]:
     """Test and end-to-end sink from the tap to the target.
 
     Note: This method buffers all output from the tap in memory and should not be
@@ -222,7 +227,7 @@ def tap_to_target_sync_test(
     target_stdout, target_stderr = target_sync_test(target, tap_stdout, finalize=True)
 
     # Reset the tap's stdout buffer before returning
-    tap_stdout.seek(0)
+    tap_stdout.rewind()
 
     return tap_stdout, tap_stderr, target_stdout, target_stderr
 
@@ -236,19 +241,23 @@ def sync_end_to_end(tap: Tap, target: Target, *mappers: InlineMapper) -> None:
         mappers: Zero or more inline mapper to apply in between the tap and target, in
             order.
     """
-    buf = io.StringIO()
-    with redirect_stdout(buf):
+    stdout_buf = RunnerStandardOutErr()
+    with redirect_stdout(stdout_buf):
         tap.sync_all()
 
-    buf.seek(0)
-    mapper_output = buf
+    # Add decoded buffer items into stdout_buf
+    stdout_buf.rewind()
+
+    mapper_output = stdout_buf
 
     for mapper in mappers:
-        buf = io.StringIO()
+        buf = RunnerStandardOutErr()
         with redirect_stdout(buf):
             mapper.listen(mapper_output)
 
-        buf.seek(0)
+        # Add decoded buffer items into stdout_buf
+        buf.rewind()
+
         mapper_output = buf
 
     target.listen(mapper_output)