Coverage for src/tomcli/toml.py: 59%
96 statements
« prev ^ index » next coverage.py v7.2.3, created at 2023-04-13 11:39 +0300
« prev ^ index » next coverage.py v7.2.3, created at 2023-04-13 11:39 +0300
1# Copyright (C) 2023 Maxwell G <maxwell@gtmx.me>
2#
3# SPDX-License-Identifier: MIT
5from __future__ import annotations
7import enum
8import io
9import sys
10from collections.abc import Iterator, Mapping, MutableMapping
11from contextlib import contextmanager
12from types import ModuleType
13from typing import IO, Any, BinaryIO
16class Reader(enum.Enum):
17 """
18 Libraries to use for deserializing TOML
19 """
21 TOMLLIB = "tomllib"
22 TOMLKIT = "tomlkit"
25class Writer(enum.Enum):
26 """
27 Libraries to use for serializing TOML
28 """
30 TOMLI_W = "tomli_w"
31 TOMLKIT = "tomlkit"
34DEFAULT_READER = Reader.TOMLKIT
35DEFAULT_WRITER = Writer.TOMLKIT
36NEEDS_STR: tuple[Writer | Reader] = [Writer.TOMLKIT]
38AVAILABLE_READERS: dict[Reader, ModuleType] = {}
39AVAILABLE_WRITERS: dict[Writer, ModuleType] = {}
41if sys.version_info[:2] >= (3, 11):
42 import tomllib
44 AVAILABLE_READERS[Reader.TOMLLIB] = tomllib
45else:
46 try:
47 import tomli as tomllib
48 except ImportError:
49 pass
50 else:
51 AVAILABLE_READERS[Reader.TOMLLIB] = tomllib
53try:
54 import tomli_w
55except ImportError:
56 pass
57else:
58 AVAILABLE_WRITERS[Writer.TOMLI_W] = tomli_w
60try:
61 import tomlkit
62except ImportError:
63 pass
64else:
65 AVAILABLE_READERS[Reader.TOMLKIT] = tomlkit
66 AVAILABLE_WRITERS[Writer.TOMLKIT] = tomlkit
69@contextmanager
70def _get_stream(fp: BinaryIO, backend: Reader | Writer) -> Iterator[IO[Any]]:
71 if backend in NEEDS_STR:
72 fp.flush()
73 wrapper = io.TextIOWrapper(fp, "utf-8")
74 try:
75 yield wrapper
76 finally:
77 wrapper.flush()
78 wrapper.detach()
79 else:
80 yield fp
83def load(
84 __fp: BinaryIO,
85 prefered_reader: Reader | None = None,
86 allow_fallback: bool = True,
87) -> MutableMapping[str, Any]:
88 """
89 Parse a bytes stream containing TOML data
91 Parameters:
92 __fp:
93 A bytes stream that supports `.read(). Positional argument only.
94 prefered_reader:
95 A [`Reader`][tomcli.toml.Reader] to use for parsing the TOML document
96 allow_fallback:
97 Whether to fallback to another Reader if `prefered_reader` is unavailable
98 """
99 prefered_reader = prefered_reader or DEFAULT_READER
100 if not AVAILABLE_READERS:
101 missing = ", ".join(module.value for module in Reader)
102 raise ModuleNotFoundError(f"None of the following were found: {missing}")
104 if prefered_reader in AVAILABLE_READERS:
105 with _get_stream(__fp, prefered_reader) as wrapper:
106 return AVAILABLE_READERS[prefered_reader].load(wrapper)
107 elif not allow_fallback:
108 raise ModuleNotFoundError(f"No module named {prefered_reader.value!r}")
110 reader, mod = next(iter(AVAILABLE_READERS.items()))
111 with _get_stream(__fp, reader) as wrapper:
112 return mod.load(wrapper)
115def dump(
116 __data: Mapping[str, Any],
117 __fp: BinaryIO,
118 prefered_writer: Writer | None = None,
119 allow_fallback: bool = True,
120) -> None:
121 """
122 Serialize an object to TOML and write it to a binary stream
124 Parameters:
125 __data:
126 A Python object to serialize. Positional argument only.
127 __fp:
128 A bytes stream that supports `.write()`. Positional argument only.
129 prefered_writer:
130 A [`Writer`][tomcli.toml.Writer] to use for serializing the Python
131 object
132 allow_fallback:
133 Whether to fallback to another Writer if `prefered_writer` is unavailable
134 """
135 prefered_writer = prefered_writer or DEFAULT_WRITER
136 if not AVAILABLE_WRITERS:
137 missing = ", ".join(module.value for module in Writer)
138 raise ModuleNotFoundError(f"None of the following were found: {missing}")
140 if prefered_writer in AVAILABLE_WRITERS:
141 with _get_stream(__fp, prefered_writer) as wrapper:
142 return AVAILABLE_WRITERS[prefered_writer].dump(__data, wrapper)
143 elif not allow_fallback:
144 raise ModuleNotFoundError(f"No module named {prefered_writer.value!r}")
146 writer, mod = next(iter(AVAILABLE_WRITERS.items()))
147 with _get_stream(__fp, writer) as wrapper:
148 return mod.dump(__data, wrapper)
150def loads(
151 __data: str,
152 prefered_reader: Reader | None = None,
153 allow_fallback: bool = True,
154) -> MutableMapping[str, Any]:
155 """
156 Parse a string containing TOML data
158 Parameters:
159 __data:
160 A string containing TOML data. Positional argument only.
161 prefered_writer:
162 A [`Writer`][tomcli.toml.Writer] to use for serializing the Python
163 object
164 allow_fallback:
165 Whether to fallback to another Writer if `prefered_writer` is unavailable
166 """
167 prefered_reader = prefered_reader or DEFAULT_READER
168 if not AVAILABLE_READERS:
169 missing = ", ".join(module.value for module in Reader)
170 raise ModuleNotFoundError(f"None of the following were found: {missing}")
172 if prefered_reader in AVAILABLE_READERS:
173 return AVAILABLE_READERS[prefered_reader].loads(__data)
174 elif not allow_fallback:
175 raise ModuleNotFoundError(f"No module named {prefered_reader.value!r}")
177 mod = next(iter(AVAILABLE_READERS.values()))
178 return mod.loads(__data)
181def dumps(
182 __data: Mapping[str, Any],
183 prefered_writer: Writer | None = None,
184 allow_fallback: bool = True,
185) -> str:
186 """
187 Serialize an object to TOML and return it as a string
189 Parameters:
190 __data:
191 A Python object to serialize. Positional argument only.
192 prefered_writer:
193 A [`Writer`][tomcli.toml.Writer] to use for serializing the Python
194 object
195 allow_fallback:
196 Whether to fallback to another Writer if `prefered_writer` is unavailable
197 """
198 prefered_writer = prefered_writer or DEFAULT_WRITER
199 if not AVAILABLE_WRITERS:
200 missing = ", ".join(module.value for module in Writer)
201 raise ModuleNotFoundError(f"None of the following were found: {missing}")
203 if prefered_writer in AVAILABLE_WRITERS:
204 return AVAILABLE_WRITERS[prefered_writer].dumps(__data)
205 elif not allow_fallback:
206 raise ModuleNotFoundError(f"No module named {prefered_writer.value!r}")
208 mod = next(iter(AVAILABLE_WRITERS.values()))
209 return mod.dumps(__data)