Coverage for src/tomcli/toml.py: 68%
72 statements
« prev ^ index » next coverage.py v7.2.3, created at 2023-04-13 04:07 +0300
« prev ^ index » next coverage.py v7.2.3, created at 2023-04-13 04:07 +0300
1# Copyright (C) 2023 Maxwell G <maxwell@gtmx.me>
2#
3# SPDX-License-Identifier: MIT
5from __future__ import annotations
6from contextlib import contextmanager
8import enum
9import io
10import sys
11from collections.abc import Iterator, Mapping, MutableMapping
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 /,
86 prefered_reader: Reader = DEFAULT_READER,
87 allow_fallback: bool = True,
88) -> MutableMapping[str, Any]:
89 """
90 Parse a bytes stream containing TOML data
92 Parameters:
93 fp:
94 A bytes stream that supports `.read()
95 prefered_reader:
96 A [`Reader`][tomcli.toml.Reader] to use for parsing the TOML document
97 allow_fallback:
98 Whether to fallback to another Reader if `prefered_reader` is unavailable
99 """
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 /,
119 prefered_writer: Writer = DEFAULT_WRITER,
120 allow_fallback: bool = True,
121) -> None:
122 """
123 Serialize an object to TOML and write it to a binary stream
125 Parameters:
126 fp:
127 A bytes stream that supports `.write()`
128 prefered_writer:
129 A [`Writer`][tomcli.toml.Writer] to use for serializing the Python
130 object
131 allow_fallback:
132 Whether to fallback to another Writer if `prefered_writer` is unavailable
133 """
134 if not AVAILABLE_WRITERS:
135 missing = ", ".join(module.value for module in Writer)
136 raise ModuleNotFoundError(f"None of the following were found: {missing}")
138 if prefered_writer in AVAILABLE_WRITERS:
139 with _get_stream(fp, prefered_writer) as wrapper:
140 return AVAILABLE_WRITERS[prefered_writer].dump(data, wrapper)
141 elif not allow_fallback:
142 raise ModuleNotFoundError(f"No module named {prefered_writer.value!r}")
144 writer, mod = next(iter(AVAILABLE_WRITERS.items()))
145 with _get_stream(fp, writer) as wrapper:
146 return mod.dump(data, wrapper)