diff --git a/dataclass_wizard/__init__.py b/dataclass_wizard/__init__.py index 97d04b18..8a7d532d 100644 --- a/dataclass_wizard/__init__.py +++ b/dataclass_wizard/__init__.py @@ -78,6 +78,7 @@ # Wizard Mixins 'JSONListWizard', 'JSONFileWizard', + 'TOMLWizard', 'YAMLWizard', # Helper serializer functions + meta config 'fromlist', @@ -104,7 +105,7 @@ Pattern, DatePattern, TimePattern, DateTimePattern) from .property_wizard import property_wizard from .serial_json import JSONSerializable -from .wizard_mixins import JSONListWizard, JSONFileWizard, YAMLWizard +from .wizard_mixins import JSONListWizard, JSONFileWizard, TOMLWizard, YAMLWizard # Set up logging to ``/dev/null`` like a library is supposed to. diff --git a/dataclass_wizard/dumpers.py b/dataclass_wizard/dumpers.py index 12ff7698..0021b9f2 100644 --- a/dataclass_wizard/dumpers.py +++ b/dataclass_wizard/dumpers.py @@ -28,7 +28,7 @@ _CLASS_TO_DUMP_FUNC, setup_dump_config_for_cls_if_needed, get_meta, dataclass_field_to_load_parser, ) -from .constants import _DUMP_HOOKS, TAG +from .constants import _DUMP_HOOKS, TAG, PY311_OR_ABOVE from .decorators import _alias from .log import LOG from .type_def import ( @@ -158,6 +158,10 @@ def setup_default_dumper(cls=DumpMixin): cls.register_dump_hook(NoneType, cls.dump_with_null) # Complex types cls.register_dump_hook(Enum, cls.dump_with_enum) + if PY311_OR_ABOVE: # Register `IntEnum` and `StrEnum` (PY 3.11+) + from enum import IntEnum, StrEnum + cls.register_dump_hook(IntEnum, cls.dump_with_enum) + cls.register_dump_hook(StrEnum, cls.dump_with_enum) cls.register_dump_hook(UUID, cls.dump_with_uuid) cls.register_dump_hook(set, cls.dump_with_iterable) cls.register_dump_hook(frozenset, cls.dump_with_iterable) diff --git a/dataclass_wizard/lazy_imports.py b/dataclass_wizard/lazy_imports.py index 8ff5a4ae..058a96d9 100644 --- a/dataclass_wizard/lazy_imports.py +++ b/dataclass_wizard/lazy_imports.py @@ -5,6 +5,7 @@ $ pip install dataclass-wizard[timedelta] """ +from .constants import PY311_OR_ABOVE from .utils.lazy_loader import LazyLoader @@ -13,3 +14,13 @@ # PyYAML: to add support for (de)serializing YAML data to dataclass instances yaml = LazyLoader(globals(), 'yaml', 'yaml', local_name='PyYAML') + +# Tomli -or- tomllib (PY 3.11+): to add support for (de)serializing TOML +# data to dataclass instances +if PY311_OR_ABOVE: + import tomllib as toml +else: + toml = LazyLoader(globals(), 'tomli', 'toml', local_name='tomli') + +# Tomli-W: to add support for serializing dataclass instances to TOML +toml_w = LazyLoader(globals(), 'tomli_w', 'toml', local_name='tomli-w') diff --git a/dataclass_wizard/loaders.py b/dataclass_wizard/loaders.py index b3f5d0f6..0c4e8160 100644 --- a/dataclass_wizard/loaders.py +++ b/dataclass_wizard/loaders.py @@ -20,7 +20,7 @@ dataclass_field_to_load_parser, json_field_to_dataclass_field, _CLASS_TO_LOAD_FUNC, dataclass_fields, get_meta, is_subclass_safe, ) -from .constants import _LOAD_HOOKS, SINGLE_ARG_ALIAS, IDENTITY +from .constants import _LOAD_HOOKS, SINGLE_ARG_ALIAS, IDENTITY, PY311_OR_ABOVE from .decorators import _alias, _single_arg_alias, resolve_alias_func, _identity from .errors import (ParseError, MissingFields, UnknownJSONKey, MissingData, RecursiveClassError) @@ -496,6 +496,10 @@ def setup_default_loader(cls=LoadMixin): cls.register_load_hook(NoneType, cls.default_load_to) # Complex types cls.register_load_hook(Enum, cls.load_to_enum) + if PY311_OR_ABOVE: # Register `IntEnum` and `StrEnum` (PY 3.11+) + from enum import IntEnum, StrEnum + cls.register_load_hook(IntEnum, cls.load_to_enum) + cls.register_load_hook(StrEnum, cls.load_to_enum) cls.register_load_hook(UUID, cls.load_to_uuid) cls.register_load_hook(set, cls.load_to_iterable) cls.register_load_hook(frozenset, cls.load_to_iterable) diff --git a/dataclass_wizard/type_def.py b/dataclass_wizard/type_def.py index a707abbd..a75d5a82 100644 --- a/dataclass_wizard/type_def.py +++ b/dataclass_wizard/type_def.py @@ -173,6 +173,8 @@ class Encoder(PyProtocol): """ def __call__(self, obj: Union[JSONObject, JSONList], + /, + *args, **kwargs) -> AnyStr: ... diff --git a/dataclass_wizard/wizard_mixins.py b/dataclass_wizard/wizard_mixins.py index c6f1b35c..9485b5e6 100644 --- a/dataclass_wizard/wizard_mixins.py +++ b/dataclass_wizard/wizard_mixins.py @@ -3,6 +3,7 @@ """ __all__ = ['JSONListWizard', 'JSONFileWizard', + 'TOMLWizard', 'YAMLWizard'] import json @@ -13,7 +14,7 @@ from .class_helper import _META from .dumpers import asdict from .enums import LetterCase -from .lazy_imports import yaml +from .lazy_imports import toml, toml_w, yaml from .loaders import fromdict, fromlist from .models import Container from .serial_json import JSONSerializable @@ -94,6 +95,113 @@ def to_json_file(self: T, file: str, mode: str = 'w', encoder(asdict(self), out_file, **encoder_kwargs) +class TOMLWizard: + # noinspection PyUnresolvedReferences + """ + A Mixin class that makes it easier to interact with TOML data. + + .. NOTE:: + By default, *NO* key transform is used in the TOML dump process. + In practice, this means that a `snake_case` field name in Python is saved + as `snake_case` to TOML; however, this can easily be customized without + the need to sub-class from :class:`JSONWizard`. + + For example: + + >>> @dataclass + >>> class MyClass(TOMLWizard, key_transform='CAMEL'): + >>> ... + + """ + def __init_subclass__(cls, key_transform=LetterCase.NONE): + """Allow easy setup of common config, such as key casing transform.""" + + # Only add the key transform if Meta config has not been specified + # for the dataclass. + if key_transform and cls not in _META: + DumpMeta(key_transform=key_transform).bind_to(cls) + + @classmethod + def from_toml(cls: Type[T], + string_or_stream: Union[AnyStr, BinaryIO], *, + decoder: Optional[Decoder] = None, + **decoder_kwargs) -> Union[T, List[T]]: + """ + Converts a TOML `string` to an instance of the dataclass, or a list of + the dataclass instances. + """ + if decoder is None: + decoder = toml.loads + + o = decoder(string_or_stream, **decoder_kwargs) + + return fromdict(cls, o) if isinstance(o, dict) else fromlist(cls, o) + + @classmethod + def from_toml_file(cls: Type[T], file: str, *, + decoder: Optional[FileDecoder] = None, + **decoder_kwargs) -> Union[T, List[T]]: + """ + Reads in the TOML file contents and converts to an instance of the + dataclass, or a list of the dataclass instances. + """ + if decoder is None: + decoder = toml.load + + with open(file, 'rb') as in_file: + return cls.from_toml(in_file, + decoder=decoder, + **decoder_kwargs) + + def to_toml(self: T, + /, + *encoder_args, + encoder: Optional[Encoder] = None, + multiline_strings: bool = False, + indent: int = 4) -> AnyStr: + """ + Converts the dataclass instance to a TOML `string` representation. + """ + if encoder is None: + encoder = toml_w.dumps + + return encoder(asdict(self), *encoder_args, + multiline_strings=multiline_strings, + indent=indent) + + def to_toml_file(self: T, file: str, mode: str = 'w', + encoder: Optional[FileEncoder] = None, + multiline_strings: bool = False, + indent: int = 4) -> None: + """ + Serializes the instance and writes it to a TOML file. + """ + if encoder is None: + encoder = toml_w.dump + + with open(file, mode) as out_file: + self.to_toml(out_file, encoder=encoder, + multiline_strings=multiline_strings, + indent=indent) + + @classmethod + def list_to_toml(cls: Type[T], + instances: List[T], + header: str = 'items', + encoder: Optional[Encoder] = None, + **encoder_kwargs) -> AnyStr: + """ + Converts a ``list`` of dataclass instances to a TOML `string` + representation. + """ + if encoder is None: + encoder = toml_w.dumps + + list_of_dict = [asdict(o, cls=cls) for o in instances] + + return encoder({header: list_of_dict}, **encoder_kwargs) + + class YAMLWizard: # noinspection PyUnresolvedReferences """ diff --git a/requirements-dev.txt b/requirements-dev.txt index 3b1dd1ce..64d8ee0b 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -13,3 +13,4 @@ watchdog[watchmedo]==6.0.0 Sphinx==7.4.7; python_version == "3.9" Sphinx==8.1.3; python_version >= "3.10" twine==5.1.1 +dataclass-wizard[toml] diff --git a/setup.py b/setup.py index 11852f56..940529a0 100644 --- a/setup.py +++ b/setup.py @@ -87,7 +87,9 @@ tests_require=test_requirements, extras_require={ 'timedelta': ['pytimeparse>=1.1.7'], - 'yaml': ['PyYAML>=5.3'], + 'toml': ['tomli>=2,<3; python_version == "3.9" or python_version == "3.10"', + 'tomli-w>=1,<2'], + 'yaml': ['PyYAML>=6,<7'], 'dev': dev_requires + doc_requires + test_requirements, }, zip_safe=False diff --git a/tests/unit/test_wizard_mixins.py b/tests/unit/test_wizard_mixins.py index 9a0b123b..9040fd38 100644 --- a/tests/unit/test_wizard_mixins.py +++ b/tests/unit/test_wizard_mixins.py @@ -1,3 +1,4 @@ +import io from dataclasses import dataclass from typing import List, Optional, Dict @@ -6,7 +7,7 @@ from dataclass_wizard import Container from dataclass_wizard.wizard_mixins import ( - JSONListWizard, JSONFileWizard, YAMLWizard + JSONListWizard, JSONFileWizard, TOMLWizard, YAMLWizard ) from .conftest import SampleClass @@ -176,3 +177,101 @@ class MyClass(YAMLWizard, key_transform=None): assert result == mock_return_val mock_encoder.assert_any_call([]) + + +@dataclass +class MyTOMLWizard(TOMLWizard): + my_str: str + inner: Optional['Inner'] = None + + +def test_toml_wizard_methods(mocker: MockerFixture): + """Test and cover the base methods in TOMLWizard.""" + toml_data = b"""\ +my_str = "test value" +[inner] +my_float = 1.2 +my_list = ["hello, world!", "123"] + """ + + # Mock open to return the TOML data as a string directly. + mock_open = mocker.patch("dataclass_wizard.wizard_mixins.open", mocker.mock_open(read_data=toml_data)) + + filename = 'my_file.toml' + + # Test reading from TOML file + obj = MyTOMLWizard.from_toml_file(filename) + + mock_open.assert_called_once_with(filename, 'rb') + mock_open.reset_mock() + + assert obj == MyTOMLWizard(my_str="test value", + inner=Inner(my_float=1.2, + my_list=["hello, world!", "123"])) + + # Test writing to TOML file + # Mock open for writing to the TOML file. + mock_open_write = mocker.mock_open() + mocker.patch("dataclass_wizard.wizard_mixins.open", mock_open_write) + + obj.to_toml_file(filename) + + mock_open_write.assert_called_once_with(filename, 'w') + + +def test_toml_wizard_list_to_toml(): + """Test and cover the `list_to_toml` method in TOMLWizard.""" + @dataclass + class MyClass(TOMLWizard, key_transform='SNAKE'): + my_str: str + my_dict: Dict[str, str] + + toml_string = MyClass.list_to_toml([ + MyClass('42', {'111': 'hello', '222': 'world'}), + MyClass('testing!', {'333': 'this is a test.'}) + ]) + + print(toml_string) + + assert toml_string == """\ +items = [ + { my_str = "42", my_dict = { 111 = "hello", 222 = "world" } }, + { my_str = "testing!", my_dict = { 333 = "this is a test." } }, +] +""" + + +def test_toml_wizard_for_branch_coverage(mocker: MockerFixture): + """Test branching logic in TOMLWizard, mainly for code coverage purposes.""" + + # This is to cover the `if` condition in the `__init_subclass__` + @dataclass + class MyClass(TOMLWizard, key_transform=None): + ... + + # from_toml: To cover the case of passing in `decoder` + mock_return_val = {'my_str': 'test string'} + + mock_decoder = mocker.Mock() + mock_decoder.return_value = mock_return_val + + result = MyTOMLWizard.from_toml('my stream', decoder=mock_decoder) + + assert result == MyTOMLWizard('test string') + mock_decoder.assert_called_once() + + # to_toml: To cover the case of passing in `encoder` + mock_encoder = mocker.Mock() + mock_encoder.return_value = mock_return_val + + m = MyTOMLWizard('test string') + result = m.to_toml(encoder=mock_encoder) + + assert result == mock_return_val + mock_encoder.assert_called_once() + + # list_to_toml: To cover the case of passing in `encoder` + result = MyTOMLWizard.list_to_toml([], encoder=mock_encoder) + + assert result == mock_return_val + mock_encoder.assert_any_call({'items': []})