comicbox.schemas.yaml

[docs] module comicbox.schemas.yaml

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
"""Comic yaml superclass."""

from collections.abc import Mapping
from decimal import Decimal
from enum import Enum
from sys import maxsize
from types import MappingProxyType
from typing import Any

from ruamel.yaml import YAML, MappingNode, RoundTripRepresenter, ScalarNode, StringIO
from typing_extensions import override

from comicbox.schemas.base import BaseRenderModule, BaseSchema, BaseSubSchema
from comicbox.schemas.comicbox import BOOKMARK_KEY, ID_KEY_KEY, PAGE_KEYS
from comicbox.schemas.comicinfo import IMAGE_ATTRIBUTE

_TAG_YAML = "tag:yaml.org,2002"
_FLOAT_TAG = f"{_TAG_YAML}:float"
_MAP_TAG = f"{_TAG_YAML}:map"
_FLOW_KEYS = frozenset({IMAGE_ATTRIBUTE, *PAGE_KEYS} - {BOOKMARK_KEY, ID_KEY_KEY})


class YamlRenderModule(BaseRenderModule):
    """Marshmallow Render Module imitates json module."""

    @staticmethod
    def _decimal_representer(dumper: RoundTripRepresenter, data: Decimal) -> ScalarNode:
        """Represent decimals as a naked 2 precision float."""
        return dumper.represent_scalar(_FLOAT_TAG, format(data, ".2f"))

    @staticmethod
    def _dict_flow_representer(
        dumper: RoundTripRepresenter, data: dict[str, Any]
    ) -> MappingNode:
        """Represent page dict as a single line."""
        if _FLOW_KEYS & data.keys():
            return dumper.represent_mapping(_MAP_TAG, data, flow_style=True)

        return dumper.represent_dict(data)

    @staticmethod
    def _none_representer(dumper: RoundTripRepresenter, data: None) -> ScalarNode:
        return dumper.represent_none(data)

    @staticmethod
    def _enum_representer(dumper: RoundTripRepresenter, data: Enum) -> ScalarNode:
        """Represent enums as their value."""
        return dumper.represent_str(data.value)

    @classmethod
    def _config_yaml(cls, yaml: YAML) -> None:
        yaml.sort_base_mapping_type_on_output = True  # pyright: ignore[reportAttributeAccessIssue]
        yaml.representer.add_representer(Decimal, cls._decimal_representer)
        yaml.representer.add_representer(type(None), cls._none_representer)
        yaml.representer.add_representer(dict, cls._dict_flow_representer)
        yaml.representer.add_multi_representer(Enum, cls._enum_representer)

    @classmethod
    def _get_write_yaml_dfs(cls) -> YAML:
        """Get write yaml with special formatting in default flow style."""
        yaml = YAML()
        yaml.default_flow_style = True
        yaml.width = maxsize
        return yaml

    @classmethod
    def _get_write_yaml(cls) -> YAML:
        """Get write yaml with special formatting."""
        yaml = YAML()
        yaml.indent(mapping=2, sequence=4, offset=2)
        return yaml

    @override
    @classmethod
    def dumps(cls, obj: Mapping, *args: Any, dfs: bool = False, **kwargs: Any) -> str:
        """Dump dict to YAML string."""
        yaml = cls._get_write_yaml_dfs() if dfs else cls._get_write_yaml()
        cls._config_yaml(yaml)
        with StringIO() as buf:
            yaml.dump(dict(obj), buf, *args, **kwargs)
            return buf.getvalue()

    @override
    @classmethod
    def loads(cls, s: str | bytes | bytearray, *args: Any, **kwargs: Any) -> Any:
        """Load YAML string into a dict."""
        if cleaned_s := cls.clean_string(s):
            return YAML().load(cleaned_s, *args, **kwargs)
        return None


class YamlSubSchema(BaseSubSchema):
    """YAML sub schema."""


class YamlSchema(BaseSchema):
    """YAML schema."""

    class Meta(BaseSchema.Meta):
        """Schema Options."""

        render_module = YamlRenderModule

    @override
    def dumps(
        self,
        obj: dict[str, Any] | MappingProxyType[str, Any],
        *args: Any,
        dfs: bool = False,
        dump: bool = True,
        **kwargs: Any,
    ) -> str:
        """Use dfs for render."""
        serialized: Any = self.dump(obj, *args, **kwargs) if dump else obj
        return self.opts.render_module.dumps(serialized, *args, dfs=dfs, **kwargs)