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

1# Copyright (C) 2023 Maxwell G <maxwell@gtmx.me> 

2# 

3# SPDX-License-Identifier: MIT 

4 

5from __future__ import annotations 

6 

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 

14 

15 

16class Reader(enum.Enum): 

17 """ 

18 Libraries to use for deserializing TOML 

19 """ 

20 

21 TOMLLIB = "tomllib" 

22 TOMLKIT = "tomlkit" 

23 

24 

25class Writer(enum.Enum): 

26 """ 

27 Libraries to use for serializing TOML 

28 """ 

29 

30 TOMLI_W = "tomli_w" 

31 TOMLKIT = "tomlkit" 

32 

33 

34DEFAULT_READER = Reader.TOMLKIT 

35DEFAULT_WRITER = Writer.TOMLKIT 

36NEEDS_STR: tuple[Writer | Reader] = [Writer.TOMLKIT] 

37 

38AVAILABLE_READERS: dict[Reader, ModuleType] = {} 

39AVAILABLE_WRITERS: dict[Writer, ModuleType] = {} 

40 

41if sys.version_info[:2] >= (3, 11): 

42 import tomllib 

43 

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 

52 

53try: 

54 import tomli_w 

55except ImportError: 

56 pass 

57else: 

58 AVAILABLE_WRITERS[Writer.TOMLI_W] = tomli_w 

59 

60try: 

61 import tomlkit 

62except ImportError: 

63 pass 

64else: 

65 AVAILABLE_READERS[Reader.TOMLKIT] = tomlkit 

66 AVAILABLE_WRITERS[Writer.TOMLKIT] = tomlkit 

67 

68 

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 

81 

82 

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 

90 

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}") 

103 

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}") 

109 

110 reader, mod = next(iter(AVAILABLE_READERS.items())) 

111 with _get_stream(__fp, reader) as wrapper: 

112 return mod.load(wrapper) 

113 

114 

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 

123 

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}") 

139 

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}") 

145 

146 writer, mod = next(iter(AVAILABLE_WRITERS.items())) 

147 with _get_stream(__fp, writer) as wrapper: 

148 return mod.dump(__data, wrapper) 

149 

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 

157 

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}") 

171 

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}") 

176 

177 mod = next(iter(AVAILABLE_READERS.values())) 

178 return mod.loads(__data) 

179 

180 

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 

188 

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}") 

202 

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}") 

207 

208 mod = next(iter(AVAILABLE_WRITERS.values())) 

209 return mod.dumps(__data)