Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.10/site-packages/tomli_w/_writer.py: 72%

1from __future__ import annotations 

2 

3from collections.abc import Generator, Mapping 

4from datetime import date, datetime, time 

5from decimal import Decimal 

6import string 

7from types import MappingProxyType 

8from typing import IO, Any, NamedTuple 

9 

10ASCII_CTRL = frozenset(chr(i) for i in range(32)) | frozenset(chr(127)) 

11ILLEGAL_BASIC_STR_CHARS = frozenset('"\\') | ASCII_CTRL - frozenset("\t") 

12BARE_KEY_CHARS = frozenset(string.ascii_letters + string.digits + "-_") 

13ARRAY_TYPES = (list, tuple) 

14MAX_LINE_LENGTH = 100 

15 

16COMPACT_ESCAPES = MappingProxyType( 

17 { 

18 "\u0008": "\\b", # backspace 

19 "\u000A": "\\n", # linefeed 

20 "\u000C": "\\f", # form feed 

21 "\u000D": "\\r", # carriage return 

22 "\u0022": '\\"', # quote 

23 "\u005C": "\\\\", # backslash 

24 } 

25) 

26 

27 

28class Context(NamedTuple): 

29 allow_multiline: bool 

30 # cache rendered inline tables (mapping from object id to rendered inline table) 

31 inline_table_cache: dict[int, str] 

32 indent_str: str 

33 

34 

35def make_context(multiline_strings: bool, indent: int) -> Context: 

36 if indent < 0: 

37 raise ValueError("Indent width must be non-negative") 

38 return Context(multiline_strings, {}, " " * indent) 

39 

40 

41def dump( 

42 obj: Mapping[str, Any], 

43 fp: IO[bytes], 

44 /, 

45 *, 

46 multiline_strings: bool = False, 

47 indent: int = 4, 

48) -> None: 

49 ctx = make_context(multiline_strings, indent) 

50 for chunk in gen_table_chunks(obj, ctx, name=""): 

51 fp.write(chunk.encode()) 

52 

53 

54def dumps( 

55 obj: Mapping[str, Any], /, *, multiline_strings: bool = False, indent: int = 4 

56) -> str: 

57 ctx = make_context(multiline_strings, indent) 

58 return "".join(gen_table_chunks(obj, ctx, name="")) 

59 

60 

61def gen_table_chunks( 

62 table: Mapping[str, Any], 

63 ctx: Context, 

64 *, 

65 name: str, 

66 inside_aot: bool = False, 

67) -> Generator[str, None, None]: 

68 yielded = False 

69 literals = [] 

70 tables: list[tuple[str, Any, bool]] = [] # => [(key, value, inside_aot)] 

71 for k, v in table.items(): 

72 if isinstance(v, Mapping): 

73 tables.append((k, v, False)) 

74 elif is_aot(v) and not all(is_suitable_inline_table(t, ctx) for t in v): 

75 tables.extend((k, t, True) for t in v) 

76 else: 

77 literals.append((k, v)) 

78 

79 if inside_aot or name and (literals or not tables): 

80 yielded = True 

81 yield f"[[{name}]]\n" if inside_aot else f"[{name}]\n" 

82 

83 if literals: 

84 yielded = True 

85 for k, v in literals: 

86 yield f"{format_key_part(k)} = {format_literal(v, ctx)}\n" 

87 

88 for k, v, in_aot in tables: 

89 if yielded: 

90 yield "\n" 

91 else: 

92 yielded = True 

93 key_part = format_key_part(k) 

94 display_name = f"{name}.{key_part}" if name else key_part 

95 yield from gen_table_chunks(v, ctx, name=display_name, inside_aot=in_aot) 

96 

97 

98def format_literal(obj: object, ctx: Context, *, nest_level: int = 0) -> str: 

99 if isinstance(obj, bool): 

100 return "true" if obj else "false" 

101 if isinstance(obj, (int, float, date, datetime)): 

102 return str(obj) 

103 if isinstance(obj, Decimal): 

104 return format_decimal(obj) 

105 if isinstance(obj, time): 

106 if obj.tzinfo: 

107 raise ValueError("TOML does not support offset times") 

108 return str(obj) 

109 if isinstance(obj, str): 

110 return format_string(obj, allow_multiline=ctx.allow_multiline) 

111 if isinstance(obj, ARRAY_TYPES): 

112 return format_inline_array(obj, ctx, nest_level) 

113 if isinstance(obj, Mapping): 

114 return format_inline_table(obj, ctx) 

115 raise TypeError( 

116 f"Object of type '{type(obj).__qualname__}' is not TOML serializable" 

117 ) 

118 

119 

120def format_decimal(obj: Decimal) -> str: 

121 if obj.is_nan(): 

122 return "nan" 

123 if obj == Decimal("inf"): 

124 return "inf" 

125 if obj == Decimal("-inf"): 

126 return "-inf" 

127 return str(obj) 

128 

129 

130def format_inline_table(obj: Mapping, ctx: Context) -> str: 

131 # check cache first 

132 obj_id = id(obj) 

133 if obj_id in ctx.inline_table_cache: 

134 return ctx.inline_table_cache[obj_id] 

135 

136 if not obj: 

137 rendered = "{}" 

138 else: 

139 rendered = ( 

140 "{ " 

141 + ", ".join( 

142 f"{format_key_part(k)} = {format_literal(v, ctx)}" 

143 for k, v in obj.items() 

144 ) 

145 + " }" 

146 ) 

147 ctx.inline_table_cache[obj_id] = rendered 

148 return rendered 

149 

150 

151def format_inline_array(obj: tuple | list, ctx: Context, nest_level: int) -> str: 

152 if not obj: 

153 return "[]" 

154 item_indent = ctx.indent_str * (1 + nest_level) 

155 closing_bracket_indent = ctx.indent_str * nest_level 

156 return ( 

157 "[\n" 

158 + ",\n".join( 

159 item_indent + format_literal(item, ctx, nest_level=nest_level + 1) 

160 for item in obj 

161 ) 

162 + f",\n{closing_bracket_indent}]" 

163 ) 

164 

165 

166def format_key_part(part: str) -> str: 

167 try: 

168 only_bare_key_chars = BARE_KEY_CHARS.issuperset(part) 

169 except TypeError: 

170 raise TypeError( 

171 f"Invalid mapping key '{part}' of type '{type(part).__qualname__}'." 

172 " A string is required." 

173 ) from None 

174 

175 if part and only_bare_key_chars: 

176 return part 

177 return format_string(part, allow_multiline=False) 

178 

179 

180def format_string(s: str, *, allow_multiline: bool) -> str: 

181 do_multiline = allow_multiline and "\n" in s 

182 if do_multiline: 

183 result = '"""\n' 

184 s = s.replace("\r\n", "\n") 

185 else: 

186 result = '"' 

187 

188 pos = seq_start = 0 

189 while True: 

190 try: 

191 char = s[pos] 

192 except IndexError: 

193 result += s[seq_start:pos] 

194 if do_multiline: 

195 return result + '"""' 

196 return result + '"' 

197 if char in ILLEGAL_BASIC_STR_CHARS: 

198 result += s[seq_start:pos] 

199 if char in COMPACT_ESCAPES: 

200 if do_multiline and char == "\n": 

201 result += "\n" 

202 else: 

203 result += COMPACT_ESCAPES[char] 

204 else: 

205 result += "\\u" + hex(ord(char))[2:].rjust(4, "0") 

206 seq_start = pos + 1 

207 pos += 1 

208 

209 

210def is_aot(obj: Any) -> bool: 

211 """Decides if an object behaves as an array of tables (i.e. a nonempty list 

212 of dicts).""" 

213 return bool( 

214 isinstance(obj, ARRAY_TYPES) 

215 and obj 

216 and all(isinstance(v, Mapping) for v in obj) 

217 ) 

218 

219 

220def is_suitable_inline_table(obj: Mapping, ctx: Context) -> bool: 

221 """Use heuristics to decide if the inline-style representation is a good 

222 choice for a given table.""" 

223 rendered_inline = f"{ctx.indent_str}{format_inline_table(obj, ctx)}," 

224 return len(rendered_inline) <= MAX_LINE_LENGTH and "\n" not in rendered_inline