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

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

2# 

3# SPDX-License-Identifier: MIT 

4 

5from __future__ import annotations 

6from contextlib import contextmanager 

7 

8import enum 

9import io 

10import sys 

11from collections.abc import Iterator, Mapping, MutableMapping 

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 /, 

86 prefered_reader: Reader = DEFAULT_READER, 

87 allow_fallback: bool = True, 

88) -> MutableMapping[str, Any]: 

89 """ 

90 Parse a bytes stream containing TOML data 

91 

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

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 /, 

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 

124 

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

137 

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

143 

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

145 with _get_stream(fp, writer) as wrapper: 

146 return mod.dump(data, wrapper)