Repository: jcrist/msgspec
Pull request #268 (2023-01-18 07:07:13)

Add TOML support
This adds TOML support in a new `msgspec.toml` module. This uses:

- `tomli` (or `tomllib` for Python 3.11+) for decoding
- `tomli_w` for encoding

Both are optional dependencies that may be installed using the new
`toml` extra.

As with YAML support, only `encode`/`decode` methods are supported for
now, no `Encoder`/`Decoder` classes.


Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 006631e..d778cf1 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -27,7 +27,7 @@ jobs:
 
       - name: Install Dependencies
         run: |
-          pip install pytest mypy pyright msgpack pyyaml coverage
+          pip install pytest mypy pyright msgpack pyyaml tomli tomli_w coverage
 
       - name: Build Msgspec
         run: |
@@ -73,7 +73,7 @@ jobs:
         os: [ubuntu-20.04, macos-11, windows-2019]
 
     env:
-      CIBW_TEST_REQUIRES: "pytest msgpack pyyaml"
+      CIBW_TEST_REQUIRES: "pytest msgpack pyyaml tomli tomli_w"
       CIBW_TEST_COMMAND: "pytest {project}/tests"
       CIBW_BUILD: "cp38-* cp39-* cp310-* cp311-*"
       CIBW_SKIP: "*-win32 *_i686 *_s390x *_ppc64le"
diff --git a/docs/source/api.rst b/docs/source/api.rst
index a1a7648..2b4321a 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -77,6 +77,16 @@ YAML
 .. autofunction:: decode
 
 
+TOML
+----
+
+.. currentmodule:: msgspec.toml
+
+.. autofunction:: encode
+
+.. autofunction:: decode
+
+
 JSON Schema
 -----------
 
diff --git a/msgspec/__init__.py b/msgspec/__init__.py
index 303857e..da81b2c 100644
--- a/msgspec/__init__.py
+++ b/msgspec/__init__.py
@@ -14,6 +14,7 @@ from ._core import (
 from . import msgpack
 from . import json
 from . import yaml
+from . import toml
 from . import inspect
 
 from ._version import get_versions
diff --git a/msgspec/__init__.pyi b/msgspec/__init__.pyi
index 9dd211e..9d97de0 100644
--- a/msgspec/__init__.pyi
+++ b/msgspec/__init__.pyi
@@ -171,6 +171,7 @@ class ValidationError(DecodeError): ...
 from . import msgpack
 from . import json
 from . import yaml
+from . import toml
 from . import inspect
 
 __version__: str
diff --git a/msgspec/toml.py b/msgspec/toml.py
new file mode 100644
index 0000000..54268c3
--- /dev/null
+++ b/msgspec/toml.py
@@ -0,0 +1,176 @@
+import datetime as _datetime
+from typing import (
+    Any,
+    Type,
+    TypeVar,
+    Optional,
+    Union,
+    Callable,
+    overload,
+)
+
+from . import (
+    to_builtins as _to_builtins,
+    from_builtins as _from_builtins,
+    DecodeError as _DecodeError,
+)
+
+__all__ = ("encode", "decode")
+
+
+def __dir__():
+    return __all__
+
+
+def _import_tomllib():
+    try:
+        import tomllib
+
+        return tomllib
+    except ImportError:
+        pass
+
+    try:
+        import tomli
+
+        return tomli
+    except ImportError:
+        raise ImportError(
+            "`msgspec.toml.decode` requires `tomli` be installed.\n\n"
+            "Please either `pip` or `conda` install it as follows:\n\n"
+            "  $ python -m pip install tomli   # using pip\n"
+            "  $ conda install tomli           # or using conda"
+        ) from None
+
+
+def _import_tomli_w():
+    try:
+        import tomli_w
+
+        return tomli_w
+    except ImportError:
+        raise ImportError(
+            "`msgspec.toml.encode` requires `tomli_w` be installed.\n\n"
+            "Please either `pip` or `conda` install it as follows:\n\n"
+            "  $ python -m pip install tomli_w   # using pip\n"
+            "  $ conda install tomli_w           # or using conda"
+        ) from None
+
+
+def encode(obj: Any, *, enc_hook: Optional[Callable[[Any], Any]] = None) -> bytes:
+    """Serialize an object as TOML.
+
+    Parameters
+    ----------
+    obj : Any
+        The object to serialize.
+    enc_hook : callable, optional
+        A callable to call for objects that aren't supported msgspec types.
+        Takes the unsupported object and should return a supported object, or
+        raise a TypeError.
+
+    Returns
+    -------
+    data : bytes
+        The serialized object.
+
+    See Also
+    --------
+    decode
+    """
+    toml = _import_tomli_w()
+    return toml.dumps(
+        _to_builtins(
+            obj,
+            builtin_types=(_datetime.datetime, _datetime.date, _datetime.time),
+            str_keys=True,
+            enc_hook=enc_hook,
+        )
+    ).encode("utf-8")
+
+
+T = TypeVar("T")
+
+
+@overload
+def decode(
+    buf: Union[bytes, str], *, dec_hook: Optional[Callable[[Type, Any], Any]] = None
+) -> Any:
+    pass
+
+
+@overload
+def decode(
+    buf: Union[bytes, str],
+    *,
+    type: Type[T] = ...,
+    dec_hook: Optional[Callable[[Type, Any], Any]] = None,
+) -> T:
+    pass
+
+
+@overload
+def decode(
+    buf: Union[bytes, str],
+    *,
+    type: Any = ...,
+    dec_hook: Optional[Callable[[Type, Any], Any]] = None,
+) -> Any:
+    pass
+
+
+def decode(
+    buf: Union[bytes, str],
+    *,
+    type: Type[T] = Any,
+    dec_hook: Optional[Callable[[Type, Any], Any]] = None,
+) -> T:
+    """Deserialize an object from TOML.
+
+    Parameters
+    ----------
+    buf : bytes-like or str
+        The message to decode.
+    type : type, optional
+        A Python type (in type annotation form) to decode the object as. If
+        provided, the message will be type checked and decoded as the specified
+        type. Defaults to `Any`, in which case the message will be decoded
+        using the default TOML types.
+    dec_hook : callable, optional
+        An optional callback for handling decoding custom types. Should have
+        the signature ``dec_hook(type: Type, obj: Any) -> Any``, where ``type``
+        is the expected message type, and ``obj`` is the decoded representation
+        composed of only basic TOML types. This hook should transform ``obj``
+        into type ``type``, or raise a ``TypeError`` if unsupported.
+
+    Returns
+    -------
+    obj : Any
+        The deserialized object.
+
+    See Also
+    --------
+    encode
+    """
+    toml = _import_tomllib()
+    if isinstance(buf, str):
+        str_buf = buf
+    elif isinstance(buf, (bytes, bytearray)):
+        str_buf = buf.decode("utf-8")
+    else:
+        # call `memoryview` first, since `bytes(1)` is actually valid
+        str_buf = bytes(memoryview(buf)).decode("utf-8")
+    try:
+        obj = toml.loads(str_buf)
+    except toml.TOMLDecodeError as exc:
+        raise _DecodeError(str(exc)) from None
+
+    if type is Any:
+        return obj
+    return _from_builtins(
+        obj,
+        type,
+        builtin_types=(_datetime.datetime, _datetime.date, _datetime.time),
+        str_keys=True,
+        dec_hook=dec_hook,
+    )
diff --git a/setup.py b/setup.py
index 175573f..d92be11 100644
--- a/setup.py
+++ b/setup.py
@@ -52,6 +52,7 @@ setup(
     ],
     extras_require={
         "yaml": ["pyyaml"],
+        "toml": ['tomli ; python_version < "3.11"', "tomli_w"],
     },
     license="BSD",
     packages=["msgspec"],


Test patch:
diff --git a/tests/basic_typing_examples.py b/tests/basic_typing_examples.py
index 485b5e7..49b0dae 100644
--- a/tests/basic_typing_examples.py
+++ b/tests/basic_typing_examples.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import datetime
 import pickle
-from typing import List, Any, Type, Union
+from typing import List, Dict, Any, Type, Union
 import msgspec
 
 def check___version__() -> None:
@@ -679,6 +679,42 @@ def check_yaml_decode_dec_hook() -> None:
 
     msgspec.yaml.decode(b"test", dec_hook=dec_hook)
 
+##########################################################
+# TOML                                                   #
+##########################################################
+
+def check_toml_encode() -> None:
+    b = msgspec.toml.encode({"a": 1})
+
+    reveal_type(b)  # assert "bytes" in typ
+
+
+def check_toml_decode_any() -> None:
+    o = msgspec.toml.decode(b"a = 1")
+    reveal_type(o)  # assert "Any" in typ
+
+
+def check_toml_decode_typed() -> None:
+    o = msgspec.toml.decode(b"a = 1", type=Dict[str, int])
+    reveal_type(o)  # assert "dict" in typ.lower() and "int" in typ
+
+
+def check_toml_decode_from_str() -> None:
+    msgspec.toml.decode("a = 1")
+    o = msgspec.toml.decode("a = 1", type=Dict[str, int])
+    reveal_type(o)  # assert "dict" in typ.lower() and "int" in typ
+
+
+def check_toml_encode_enc_hook() -> None:
+    msgspec.toml.encode(object(), enc_hook=lambda x: None)
+
+
+def check_toml_decode_dec_hook() -> None:
+    def dec_hook(typ: Type, obj: Any) -> Any:
+        return typ(obj)
+
+    msgspec.toml.decode(b"a = 1", dec_hook=dec_hook)
+
 
 ##########################################################
 # msgspec.inspect                                        #
diff --git a/tests/test_toml.py b/tests/test_toml.py
new file mode 100644
index 0000000..f560677
--- /dev/null
+++ b/tests/test_toml.py
@@ -0,0 +1,190 @@
+import dataclasses
+import datetime
+import enum
+import sys
+import uuid
+from decimal import Decimal
+from typing import Dict, List, Tuple, Set, FrozenSet
+
+import pytest
+
+import msgspec
+
+try:
+    import tomllib
+except ImportError:
+    try:
+        import tomli as tomllib
+    except ImportError:
+        tomllib = None
+
+try:
+    import tomli_w
+except ImportError:
+    tomli_w = None
+
+
+needs_decode = pytest.mark.skipif(
+    tomllib is None, reason="Neither tomllib or tomli are installed"
+)
+needs_encode = pytest.mark.skipif(tomli_w is None, reason="tomli_w is not installed")
+
+PY311 = sys.version_info[:2] >= (3, 11)
+
+UTC = datetime.timezone.utc
+
+
+class ExStruct(msgspec.Struct):
+    x: int
+    y: str
+
+
+@dataclasses.dataclass
+class ExDataclass:
+    x: int
+    y: str
+
+
+class ExEnum(enum.Enum):
+    one = "one"
+    two = "two"
+
+
+class ExIntEnum(enum.IntEnum):
+    one = 1
+    two = 2
+
+
+def test_module_dir():
+    assert set(dir(msgspec.toml)) == {"encode", "decode"}
+
+
+@pytest.mark.skipif(PY311, reason="tomllib is builtin in 3.11+")
+def test_tomli_not_installed_error(monkeypatch):
+    monkeypatch.setitem(sys.modules, "tomli", None)
+
+    with pytest.raises(ImportError, match="conda install"):
+        msgspec.toml.decode("a = 1", type=int)
+
+
+def test_tomli_w_not_installed_error(monkeypatch):
+    monkeypatch.setitem(sys.modules, "tomli_w", None)
+
+    with pytest.raises(ImportError, match="conda install"):
+        msgspec.toml.encode({"a": 1})
+
+
+@pytest.mark.parametrize(
+    "val",
+    [
+        True,
+        False,
+        1,
+        1.5,
+        "fizz",
+        datetime.datetime(2022, 1, 2, 3, 4, 5, 6),
+        datetime.datetime(2022, 1, 2, 3, 4, 5, 6, UTC),
+        datetime.date(2022, 1, 2),
+        datetime.time(12, 34),
+        [1, 2],
+        {"one": 2},
+    ],
+)
+@needs_encode
+@needs_decode
+def test_roundtrip_any(val):
+    msg = msgspec.toml.encode({"x": val})
+    res = msgspec.toml.decode(msg)["x"]
+    assert res == val
+
+
+@pytest.mark.parametrize(
+    "val, type",
+    [
+        (True, bool),
+        (False, bool),
+        (1, int),
+        (1.5, float),
+        ("fizz", str),
+        (b"fizz", bytes),
+        (b"fizz", bytearray),
+        (datetime.datetime(2022, 1, 2, 3, 4, 5, 6), datetime.datetime),
+        (datetime.datetime(2022, 1, 2, 3, 4, 5, 6, UTC), datetime.datetime),
+        (datetime.date(2022, 1, 2), datetime.date),
+        (datetime.time(12, 34), datetime.time),
+        (uuid.uuid4(), uuid.UUID),
+        (ExEnum.one, ExEnum),
+        (ExIntEnum.one, ExIntEnum),
+        ([1, 2], List[int]),
+        ((1, 2), Tuple[int, ...]),
+        ({1, 2}, Set[int]),
+        (frozenset({1, 2}), FrozenSet[int]),
+        (("one", 2), Tuple[str, int]),
+        ({"one": 2}, Dict[str, int]),
+        ({1: "two"}, Dict[int, str]),
+        (ExStruct(1, "two"), ExStruct),
+        (ExDataclass(1, "two"), ExDataclass),
+    ],
+)
+@needs_encode
+@needs_decode
+def test_roundtrip_typed(val, type):
+    msg = msgspec.toml.encode({"x": val})
+    res = msgspec.toml.decode(msg, type=Dict[str, type])["x"]
+    assert res == val
+
+
+@needs_encode
+def test_encode_output_type():
+    msg = msgspec.toml.encode({"x": 1})
+    assert isinstance(msg, bytes)
+
+
+@needs_encode
+def test_encode_error():
+    class Oops:
+        pass
+
+    with pytest.raises(TypeError, match="Encoding objects of type Oops is unsupported"):
+        msgspec.toml.encode({"x": Oops()})
+
+
+@needs_encode
+@needs_decode
+def test_encode_enc_hook():
+    msg = msgspec.toml.encode({"x": Decimal(1.5)}, enc_hook=str)
+    assert msgspec.toml.decode(msg) == {"x": "1.5"}
+
+
+@needs_decode
+def test_decode_str_or_bytes_like():
+    assert msgspec.toml.decode("a = 1") == {"a": 1}
+    assert msgspec.toml.decode(b"a = 1") == {"a": 1}
+    assert msgspec.toml.decode(bytearray(b"a = 1")) == {"a": 1}
+    assert msgspec.toml.decode(memoryview(b"a = 1")) == {"a": 1}
+    with pytest.raises(TypeError):
+        msgspec.toml.decode(1)
+
+
+@needs_decode
+@pytest.mark.parametrize("msg", [b"{{", b"!!binary 123"])
+def test_decode_parse_error(msg):
+    with pytest.raises(msgspec.DecodeError):
+        msgspec.toml.decode(msg)
+
+
+@needs_decode
+def test_decode_validation_error():
+    with pytest.raises(msgspec.ValidationError, match="Expected `str`"):
+        msgspec.toml.decode(b"a = [1, 2, 3]", type=Dict[str, List[str]])
+
+
+@needs_decode
+def test_decode_dec_hook():
+    def dec_hook(typ, val):
+        if typ is Decimal:
+            return Decimal(val)
+        raise TypeError
+
+    res = msgspec.toml.decode("a = '1.5'", type=Dict[str, Decimal], dec_hook=dec_hook)
+    assert res == {"a": Decimal("1.5")}