container.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953
  1. from __future__ import annotations
  2. import copy
  3. from collections.abc import Iterator
  4. from typing import Any
  5. from tomlkit._compat import decode
  6. from tomlkit._types import _CustomDict
  7. from tomlkit._utils import merge_dicts
  8. from tomlkit.exceptions import KeyAlreadyPresent
  9. from tomlkit.exceptions import NonExistentKey
  10. from tomlkit.exceptions import TOMLKitError
  11. from tomlkit.items import AoT
  12. from tomlkit.items import Comment
  13. from tomlkit.items import Item
  14. from tomlkit.items import Key
  15. from tomlkit.items import Null
  16. from tomlkit.items import SingleKey
  17. from tomlkit.items import Table
  18. from tomlkit.items import Trivia
  19. from tomlkit.items import Whitespace
  20. from tomlkit.items import item as _item
  21. _NOT_SET = object()
  22. class Container(_CustomDict):
  23. """
  24. A container for items within a TOMLDocument.
  25. This class implements the `dict` interface with copy/deepcopy protocol.
  26. """
  27. def __init__(self, parsed: bool = False) -> None:
  28. self._map: dict[SingleKey, int | tuple[int, ...]] = {}
  29. self._body: list[tuple[Key | None, Item]] = []
  30. self._parsed = parsed
  31. self._table_keys = []
  32. @property
  33. def body(self) -> list[tuple[Key | None, Item]]:
  34. return self._body
  35. def unwrap(self) -> dict[str, Any]:
  36. """Returns as pure python object (ppo)"""
  37. unwrapped = {}
  38. for k, v in self.items():
  39. if k is None:
  40. continue
  41. if isinstance(k, Key):
  42. k = k.key
  43. if hasattr(v, "unwrap"):
  44. v = v.unwrap()
  45. if k in unwrapped:
  46. merge_dicts(unwrapped[k], v)
  47. else:
  48. unwrapped[k] = v
  49. return unwrapped
  50. @property
  51. def value(self) -> dict[str, Any]:
  52. """The wrapped dict value"""
  53. d = {}
  54. for k, v in self._body:
  55. if k is None:
  56. continue
  57. k = k.key
  58. v = v.value
  59. if isinstance(v, Container):
  60. v = v.value
  61. if k in d:
  62. merge_dicts(d[k], v)
  63. else:
  64. d[k] = v
  65. return d
  66. def parsing(self, parsing: bool) -> None:
  67. self._parsed = parsing
  68. for _, v in self._body:
  69. if isinstance(v, Table):
  70. v.value.parsing(parsing)
  71. elif isinstance(v, AoT):
  72. for t in v.body:
  73. t.value.parsing(parsing)
  74. def add(self, key: Key | Item | str, item: Item | None = None) -> Container:
  75. """
  76. Adds an item to the current Container.
  77. :Example:
  78. >>> # add a key-value pair
  79. >>> doc.add('key', 'value')
  80. >>> # add a comment or whitespace or newline
  81. >>> doc.add(comment('# comment'))
  82. """
  83. if item is None:
  84. if not isinstance(key, (Comment, Whitespace)):
  85. raise ValueError(
  86. "Non comment/whitespace items must have an associated key"
  87. )
  88. key, item = None, key
  89. return self.append(key, item)
  90. def _handle_dotted_key(self, key: Key, value: Item) -> None:
  91. if isinstance(value, (Table, AoT)):
  92. raise TOMLKitError("Can't add a table to a dotted key")
  93. name, *mid, last = key
  94. name._dotted = True
  95. table = current = Table(Container(True), Trivia(), False, is_super_table=True)
  96. for _name in mid:
  97. _name._dotted = True
  98. new_table = Table(Container(True), Trivia(), False, is_super_table=True)
  99. current.append(_name, new_table)
  100. current = new_table
  101. last.sep = key.sep
  102. current.append(last, value)
  103. self.append(name, table)
  104. return
  105. def _get_last_index_before_table(self) -> int:
  106. last_index = -1
  107. for i, (k, v) in enumerate(self._body):
  108. if isinstance(v, Null):
  109. continue # Null elements are inserted after deletion
  110. if isinstance(v, Whitespace) and not v.is_fixed():
  111. continue
  112. if isinstance(v, (Table, AoT)) and not k.is_dotted():
  113. break
  114. last_index = i
  115. return last_index + 1
  116. def _validate_out_of_order_table(self, key: SingleKey | None = None) -> None:
  117. if key is None:
  118. for k in self._map:
  119. assert k is not None
  120. self._validate_out_of_order_table(k)
  121. return
  122. if key not in self._map or not isinstance(self._map[key], tuple):
  123. return
  124. OutOfOrderTableProxy.validate(self, self._map[key])
  125. def append(
  126. self, key: Key | str | None, item: Item, validate: bool = True
  127. ) -> Container:
  128. """Similar to :meth:`add` but both key and value must be given."""
  129. if not isinstance(key, Key) and key is not None:
  130. key = SingleKey(key)
  131. if not isinstance(item, Item):
  132. item = _item(item)
  133. if key is not None and key.is_multi():
  134. self._handle_dotted_key(key, item)
  135. return self
  136. if isinstance(item, (AoT, Table)) and item.name is None:
  137. item.name = key.key
  138. prev = self._previous_item()
  139. prev_ws = isinstance(prev, Whitespace) or ends_with_whitespace(prev)
  140. if isinstance(item, Table):
  141. if not self._parsed:
  142. item.invalidate_display_name()
  143. if (
  144. self._body
  145. and not (self._parsed or item.trivia.indent or prev_ws)
  146. and not key.is_dotted()
  147. ):
  148. item.trivia.indent = "\n"
  149. if isinstance(item, AoT) and self._body and not self._parsed:
  150. item.invalidate_display_name()
  151. if item and not ("\n" in item[0].trivia.indent or prev_ws):
  152. item[0].trivia.indent = "\n" + item[0].trivia.indent
  153. if key is not None and key in self:
  154. current_idx = self._map[key]
  155. if isinstance(current_idx, tuple):
  156. current_body_element = self._body[current_idx[-1]]
  157. else:
  158. current_body_element = self._body[current_idx]
  159. current = current_body_element[1]
  160. if isinstance(item, Table):
  161. if not isinstance(current, (Table, AoT)):
  162. raise KeyAlreadyPresent(key)
  163. if item.is_aot_element():
  164. # New AoT element found later on
  165. # Adding it to the current AoT
  166. if not isinstance(current, AoT):
  167. current = AoT([current, item], parsed=self._parsed)
  168. self._replace(key, key, current)
  169. else:
  170. current.append(item)
  171. return self
  172. elif current.is_aot():
  173. if not item.is_aot_element():
  174. # Tried to define a table after an AoT with the same name.
  175. raise KeyAlreadyPresent(key)
  176. current.append(item)
  177. return self
  178. elif current.is_super_table():
  179. if item.is_super_table():
  180. # We need to merge both super tables
  181. if (
  182. key.is_dotted()
  183. or current_body_element[0].is_dotted()
  184. or self._table_keys[-1] != current_body_element[0]
  185. ):
  186. if key.is_dotted() and not self._parsed:
  187. idx = self._get_last_index_before_table()
  188. else:
  189. idx = len(self._body)
  190. if idx < len(self._body):
  191. self._insert_at(idx, key, item)
  192. else:
  193. self._raw_append(key, item)
  194. if validate:
  195. self._validate_out_of_order_table(key)
  196. return self
  197. # Create a new element to replace the old one
  198. current = copy.deepcopy(current)
  199. for k, v in item.value.body:
  200. current.append(k, v)
  201. self._body[
  202. (
  203. current_idx[-1]
  204. if isinstance(current_idx, tuple)
  205. else current_idx
  206. )
  207. ] = (current_body_element[0], current)
  208. return self
  209. elif current_body_element[0].is_dotted():
  210. raise TOMLKitError("Redefinition of an existing table")
  211. elif not item.is_super_table():
  212. raise KeyAlreadyPresent(key)
  213. elif isinstance(item, AoT):
  214. if not isinstance(current, AoT):
  215. # Tried to define an AoT after a table with the same name.
  216. raise KeyAlreadyPresent(key)
  217. for table in item.body:
  218. current.append(table)
  219. return self
  220. else:
  221. raise KeyAlreadyPresent(key)
  222. is_table = isinstance(item, (Table, AoT))
  223. if (
  224. key is not None
  225. and self._body
  226. and not self._parsed
  227. and (not is_table or key.is_dotted())
  228. ):
  229. # If there is already at least one table in the current container
  230. # and the given item is not a table, we need to find the last
  231. # item that is not a table and insert after it
  232. # If no such item exists, insert at the top of the table
  233. last_index = self._get_last_index_before_table()
  234. if last_index < len(self._body):
  235. after_item = self._body[last_index][1]
  236. if not (
  237. isinstance(after_item, Whitespace)
  238. or "\n" in after_item.trivia.indent
  239. ):
  240. after_item.trivia.indent = "\n" + after_item.trivia.indent
  241. return self._insert_at(last_index, key, item)
  242. else:
  243. previous_item = self._body[-1][1]
  244. if not (
  245. isinstance(previous_item, Whitespace)
  246. or ends_with_whitespace(previous_item)
  247. or "\n" in previous_item.trivia.trail
  248. ):
  249. previous_item.trivia.trail += "\n"
  250. self._raw_append(key, item)
  251. return self
  252. def _raw_append(self, key: Key | None, item: Item) -> None:
  253. if key in self._map:
  254. current_idx = self._map[key]
  255. if not isinstance(current_idx, tuple):
  256. current_idx = (current_idx,)
  257. current = self._body[current_idx[-1]][1]
  258. if key is not None and not isinstance(current, Table):
  259. raise KeyAlreadyPresent(key)
  260. self._map[key] = (*current_idx, len(self._body))
  261. elif key is not None:
  262. self._map[key] = len(self._body)
  263. self._body.append((key, item))
  264. if item.is_table():
  265. self._table_keys.append(key)
  266. if key is not None:
  267. dict.__setitem__(self, key.key, item.value)
  268. def _remove_at(self, idx: int) -> None:
  269. key = self._body[idx][0]
  270. index = self._map.get(key)
  271. if index is None:
  272. raise NonExistentKey(key)
  273. self._body[idx] = (None, Null())
  274. if isinstance(index, tuple):
  275. index = list(index)
  276. index.remove(idx)
  277. if len(index) == 1:
  278. index = index.pop()
  279. else:
  280. index = tuple(index)
  281. self._map[key] = index
  282. else:
  283. dict.__delitem__(self, key.key)
  284. self._map.pop(key)
  285. def remove(self, key: Key | str) -> Container:
  286. """Remove a key from the container."""
  287. if not isinstance(key, Key):
  288. key = SingleKey(key)
  289. idx = self._map.pop(key, None)
  290. if idx is None:
  291. raise NonExistentKey(key)
  292. if isinstance(idx, tuple):
  293. for i in idx:
  294. self._body[i] = (None, Null())
  295. else:
  296. self._body[idx] = (None, Null())
  297. dict.__delitem__(self, key.key)
  298. return self
  299. def _insert_after(
  300. self, key: Key | str, other_key: Key | str, item: Any
  301. ) -> Container:
  302. if key is None:
  303. raise ValueError("Key cannot be null in insert_after()")
  304. if key not in self:
  305. raise NonExistentKey(key)
  306. if not isinstance(key, Key):
  307. key = SingleKey(key)
  308. if not isinstance(other_key, Key):
  309. other_key = SingleKey(other_key)
  310. item = _item(item)
  311. idx = self._map[key]
  312. # Insert after the max index if there are many.
  313. if isinstance(idx, tuple):
  314. idx = max(idx)
  315. current_item = self._body[idx][1]
  316. if "\n" not in current_item.trivia.trail:
  317. current_item.trivia.trail += "\n"
  318. # Increment indices after the current index
  319. for k, v in self._map.items():
  320. if isinstance(v, tuple):
  321. new_indices = []
  322. for v_ in v:
  323. if v_ > idx:
  324. v_ = v_ + 1
  325. new_indices.append(v_)
  326. self._map[k] = tuple(new_indices)
  327. elif v > idx:
  328. self._map[k] = v + 1
  329. self._map[other_key] = idx + 1
  330. self._body.insert(idx + 1, (other_key, item))
  331. if key is not None:
  332. dict.__setitem__(self, other_key.key, item.value)
  333. return self
  334. def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container:
  335. if idx > len(self._body) - 1:
  336. raise ValueError(f"Unable to insert at position {idx}")
  337. if not isinstance(key, Key):
  338. key = SingleKey(key)
  339. item = _item(item)
  340. if idx > 0:
  341. previous_item = self._body[idx - 1][1]
  342. if not (
  343. isinstance(previous_item, Whitespace)
  344. or ends_with_whitespace(previous_item)
  345. or isinstance(item, (AoT, Table))
  346. or "\n" in previous_item.trivia.trail
  347. ):
  348. previous_item.trivia.trail += "\n"
  349. # Increment indices after the current index
  350. for k, v in self._map.items():
  351. if isinstance(v, tuple):
  352. new_indices = []
  353. for v_ in v:
  354. if v_ >= idx:
  355. v_ = v_ + 1
  356. new_indices.append(v_)
  357. self._map[k] = tuple(new_indices)
  358. elif v >= idx:
  359. self._map[k] = v + 1
  360. if key in self._map:
  361. current_idx = self._map[key]
  362. if not isinstance(current_idx, tuple):
  363. current_idx = (current_idx,)
  364. self._map[key] = (*current_idx, idx)
  365. else:
  366. self._map[key] = idx
  367. self._body.insert(idx, (key, item))
  368. dict.__setitem__(self, key.key, item.value)
  369. return self
  370. def item(self, key: Key | str) -> Item:
  371. """Get an item for the given key."""
  372. if not isinstance(key, Key):
  373. key = SingleKey(key)
  374. idx = self._map.get(key)
  375. if idx is None:
  376. raise NonExistentKey(key)
  377. if isinstance(idx, tuple):
  378. # The item we are getting is an out of order table
  379. # so we need a proxy to retrieve the proper objects
  380. # from the parent container
  381. return OutOfOrderTableProxy(self, idx)
  382. return self._body[idx][1]
  383. def last_item(self) -> Item | None:
  384. """Get the last item."""
  385. if self._body:
  386. return self._body[-1][1]
  387. def as_string(self) -> str:
  388. """Render as TOML string."""
  389. s = ""
  390. for k, v in self._body:
  391. if k is not None:
  392. if isinstance(v, Table):
  393. if (
  394. s.strip(" ")
  395. and not s.strip(" ").endswith("\n")
  396. and "\n" not in v.trivia.indent
  397. ):
  398. s += "\n"
  399. s += self._render_table(k, v)
  400. elif isinstance(v, AoT):
  401. if (
  402. s.strip(" ")
  403. and not s.strip(" ").endswith("\n")
  404. and "\n" not in v.trivia.indent
  405. ):
  406. s += "\n"
  407. s += self._render_aot(k, v)
  408. else:
  409. s += self._render_simple_item(k, v)
  410. else:
  411. s += self._render_simple_item(k, v)
  412. return s
  413. def _render_table(self, key: Key, table: Table, prefix: str | None = None) -> str:
  414. cur = ""
  415. if table.display_name is not None:
  416. _key = table.display_name
  417. else:
  418. _key = key.as_string()
  419. if prefix is not None:
  420. _key = prefix + "." + _key
  421. if (
  422. not table.is_super_table()
  423. or (
  424. any(
  425. not isinstance(v, (Table, AoT, Whitespace, Null))
  426. for _, v in table.value.body
  427. )
  428. and not key.is_dotted()
  429. )
  430. or (
  431. any(k.is_dotted() for k, v in table.value.body if isinstance(v, Table))
  432. and not key.is_dotted()
  433. )
  434. ):
  435. open_, close = "[", "]"
  436. if table.is_aot_element():
  437. open_, close = "[[", "]]"
  438. newline_in_table_trivia = (
  439. "\n" if "\n" not in table.trivia.trail and len(table.value) > 0 else ""
  440. )
  441. cur += (
  442. f"{table.trivia.indent}"
  443. f"{open_}"
  444. f"{decode(_key)}"
  445. f"{close}"
  446. f"{table.trivia.comment_ws}"
  447. f"{decode(table.trivia.comment)}"
  448. f"{table.trivia.trail}"
  449. f"{newline_in_table_trivia}"
  450. )
  451. elif table.trivia.indent == "\n":
  452. cur += table.trivia.indent
  453. for k, v in table.value.body:
  454. if isinstance(v, Table):
  455. if (
  456. cur.strip(" ")
  457. and not cur.strip(" ").endswith("\n")
  458. and "\n" not in v.trivia.indent
  459. ):
  460. cur += "\n"
  461. if v.is_super_table():
  462. if k.is_dotted() and not key.is_dotted():
  463. # Dotted key inside table
  464. cur += self._render_table(k, v)
  465. else:
  466. cur += self._render_table(k, v, prefix=_key)
  467. else:
  468. cur += self._render_table(k, v, prefix=_key)
  469. elif isinstance(v, AoT):
  470. if (
  471. cur.strip(" ")
  472. and not cur.strip(" ").endswith("\n")
  473. and "\n" not in v.trivia.indent
  474. ):
  475. cur += "\n"
  476. cur += self._render_aot(k, v, prefix=_key)
  477. else:
  478. cur += self._render_simple_item(
  479. k, v, prefix=_key if key.is_dotted() else None
  480. )
  481. return cur
  482. def _render_aot(self, key, aot, prefix=None):
  483. _key = key.as_string()
  484. if prefix is not None:
  485. _key = prefix + "." + _key
  486. cur = ""
  487. _key = decode(_key)
  488. for table in aot.body:
  489. cur += self._render_aot_table(table, prefix=_key)
  490. return cur
  491. def _render_aot_table(self, table: Table, prefix: str | None = None) -> str:
  492. cur = ""
  493. _key = prefix or ""
  494. open_, close = "[[", "]]"
  495. cur += (
  496. f"{table.trivia.indent}"
  497. f"{open_}"
  498. f"{decode(_key)}"
  499. f"{close}"
  500. f"{table.trivia.comment_ws}"
  501. f"{decode(table.trivia.comment)}"
  502. f"{table.trivia.trail}"
  503. )
  504. for k, v in table.value.body:
  505. if isinstance(v, Table):
  506. if v.is_super_table():
  507. if k.is_dotted():
  508. # Dotted key inside table
  509. cur += self._render_table(k, v)
  510. else:
  511. cur += self._render_table(k, v, prefix=_key)
  512. else:
  513. cur += self._render_table(k, v, prefix=_key)
  514. elif isinstance(v, AoT):
  515. cur += self._render_aot(k, v, prefix=_key)
  516. else:
  517. cur += self._render_simple_item(k, v)
  518. return cur
  519. def _render_simple_item(self, key, item, prefix=None):
  520. if key is None:
  521. return item.as_string()
  522. _key = key.as_string()
  523. if prefix is not None:
  524. _key = prefix + "." + _key
  525. return (
  526. f"{item.trivia.indent}"
  527. f"{decode(_key)}"
  528. f"{key.sep}"
  529. f"{decode(item.as_string())}"
  530. f"{item.trivia.comment_ws}"
  531. f"{decode(item.trivia.comment)}"
  532. f"{item.trivia.trail}"
  533. )
  534. def __len__(self) -> int:
  535. return dict.__len__(self)
  536. def __iter__(self) -> Iterator[str]:
  537. return iter(dict.keys(self))
  538. # Dictionary methods
  539. def __getitem__(self, key: Key | str) -> Item | Container:
  540. item = self.item(key)
  541. if isinstance(item, Item) and item.is_boolean():
  542. return item.value
  543. return item
  544. def __setitem__(self, key: Key | str, value: Any) -> None:
  545. if key is not None and key in self:
  546. old_key = next(filter(lambda k: k == key, self._map))
  547. self._replace(old_key, key, value)
  548. else:
  549. self.append(key, value)
  550. def __delitem__(self, key: Key | str) -> None:
  551. self.remove(key)
  552. def setdefault(self, key: Key | str, default: Any) -> Any:
  553. super().setdefault(key, default=default)
  554. return self[key]
  555. def _replace(self, key: Key | str, new_key: Key | str, value: Item) -> None:
  556. if not isinstance(key, Key):
  557. key = SingleKey(key)
  558. idx = self._map.get(key)
  559. if idx is None:
  560. raise NonExistentKey(key)
  561. self._replace_at(idx, new_key, value)
  562. def _replace_at(
  563. self, idx: int | tuple[int], new_key: Key | str, value: Item
  564. ) -> None:
  565. value = _item(value)
  566. if isinstance(idx, tuple):
  567. for i in idx[1:]:
  568. self._body[i] = (None, Null())
  569. idx = idx[0]
  570. k, v = self._body[idx]
  571. if not isinstance(new_key, Key):
  572. if (
  573. isinstance(value, (AoT, Table)) != isinstance(v, (AoT, Table))
  574. or new_key != k.key
  575. ):
  576. new_key = SingleKey(new_key)
  577. else: # Inherit the sep of the old key
  578. new_key = k
  579. del self._map[k]
  580. self._map[new_key] = idx
  581. if new_key != k:
  582. dict.__delitem__(self, k)
  583. if isinstance(value, (AoT, Table)) != isinstance(v, (AoT, Table)):
  584. # new tables should appear after all non-table values
  585. self.remove(k)
  586. for i in range(idx, len(self._body)):
  587. if isinstance(self._body[i][1], (AoT, Table)):
  588. self._insert_at(i, new_key, value)
  589. idx = i
  590. break
  591. else:
  592. idx = -1
  593. self.append(new_key, value)
  594. else:
  595. # Copying trivia
  596. if not isinstance(value, (Whitespace, AoT)):
  597. value.trivia.indent = v.trivia.indent
  598. value.trivia.comment_ws = value.trivia.comment_ws or v.trivia.comment_ws
  599. value.trivia.comment = value.trivia.comment or v.trivia.comment
  600. value.trivia.trail = v.trivia.trail
  601. self._body[idx] = (new_key, value)
  602. if hasattr(value, "invalidate_display_name"):
  603. value.invalidate_display_name() # type: ignore[attr-defined]
  604. if isinstance(value, Table):
  605. # Insert a cosmetic new line for tables if:
  606. # - it does not have it yet OR is not followed by one
  607. # - it is not the last item, or
  608. # - The table being replaced has a newline
  609. last, _ = self._previous_item_with_index()
  610. idx = last if idx < 0 else idx
  611. has_ws = ends_with_whitespace(value)
  612. replace_has_ws = (
  613. isinstance(v, Table)
  614. and v.value.body
  615. and isinstance(v.value.body[-1][1], Whitespace)
  616. )
  617. next_ws = idx < last and isinstance(self._body[idx + 1][1], Whitespace)
  618. if (idx < last or replace_has_ws) and not (next_ws or has_ws):
  619. value.append(None, Whitespace("\n"))
  620. dict.__setitem__(self, new_key.key, value.value)
  621. def __str__(self) -> str:
  622. return str(self.value)
  623. def __repr__(self) -> str:
  624. return repr(self.value)
  625. def __eq__(self, other: dict) -> bool:
  626. if not isinstance(other, dict):
  627. return NotImplemented
  628. return self.value == other
  629. def _getstate(self, protocol):
  630. return (self._parsed,)
  631. def __reduce__(self):
  632. return self.__reduce_ex__(2)
  633. def __reduce_ex__(self, protocol):
  634. return (
  635. self.__class__,
  636. self._getstate(protocol),
  637. (self._map, self._body, self._parsed, self._table_keys),
  638. )
  639. def __setstate__(self, state):
  640. self._map = state[0]
  641. self._body = state[1]
  642. self._parsed = state[2]
  643. self._table_keys = state[3]
  644. for key, item in self._body:
  645. if key is not None:
  646. dict.__setitem__(self, key.key, item.value)
  647. def copy(self) -> Container:
  648. return copy.copy(self)
  649. def __copy__(self) -> Container:
  650. c = self.__class__(self._parsed)
  651. for k, v in dict.items(self):
  652. dict.__setitem__(c, k, v)
  653. c._body += self.body
  654. c._map.update(self._map)
  655. return c
  656. def _previous_item_with_index(
  657. self, idx: int | None = None, ignore=(Null,)
  658. ) -> tuple[int, Item] | None:
  659. """Find the immediate previous item before index ``idx``"""
  660. if idx is None or idx > len(self._body):
  661. idx = len(self._body)
  662. for i in range(idx - 1, -1, -1):
  663. v = self._body[i][-1]
  664. if not isinstance(v, ignore):
  665. return i, v
  666. return None
  667. def _previous_item(self, idx: int | None = None, ignore=(Null,)) -> Item | None:
  668. """Find the immediate previous item before index ``idx``.
  669. If ``idx`` is not given, the last item is returned.
  670. """
  671. prev = self._previous_item_with_index(idx, ignore)
  672. return prev[-1] if prev else None
  673. class OutOfOrderTableProxy(_CustomDict):
  674. @staticmethod
  675. def validate(container: Container, indices: tuple[int, ...]) -> None:
  676. """Validate out of order tables in the given container"""
  677. # Append all items to a temp container to see if there is any error
  678. temp_container = Container(True)
  679. for i in indices:
  680. _, item = container._body[i]
  681. if isinstance(item, Table):
  682. for k, v in item.value.body:
  683. temp_container.append(k, v, validate=False)
  684. temp_container._validate_out_of_order_table()
  685. def __init__(self, container: Container, indices: tuple[int, ...]) -> None:
  686. self._container = container
  687. self._internal_container = Container(True)
  688. self._tables: list[Table] = []
  689. self._tables_map: dict[Key, list[int]] = {}
  690. for i in indices:
  691. _, item = self._container._body[i]
  692. if isinstance(item, Table):
  693. self._tables.append(item)
  694. table_idx = len(self._tables) - 1
  695. for k, v in item.value.body:
  696. self._internal_container._raw_append(k, v)
  697. indices = self._tables_map.setdefault(k, [])
  698. if table_idx not in indices:
  699. indices.append(table_idx)
  700. if k is not None:
  701. dict.__setitem__(self, k.key, v)
  702. self._internal_container._validate_out_of_order_table()
  703. def unwrap(self) -> str:
  704. return self._internal_container.unwrap()
  705. @property
  706. def value(self):
  707. return self._internal_container.value
  708. def __getitem__(self, key: Key | str) -> Any:
  709. if key not in self._internal_container:
  710. raise NonExistentKey(key)
  711. return self._internal_container[key]
  712. def __setitem__(self, key: Key | str, value: Any) -> None:
  713. from .items import item
  714. def _is_table_or_aot(it: Any) -> bool:
  715. return isinstance(item(it), (Table, AoT))
  716. if key in self._tables_map:
  717. # Overwrite the first table and remove others
  718. indices = self._tables_map[key]
  719. while len(indices) > 1:
  720. table = self._tables[indices.pop()]
  721. self._remove_table(table)
  722. old_value = self._tables[indices[0]][key]
  723. if _is_table_or_aot(old_value) and not _is_table_or_aot(value):
  724. # Remove the entry from the map and set value again.
  725. del self._tables[indices[0]][key]
  726. del self._tables_map[key]
  727. self[key] = value
  728. return
  729. self._tables[indices[0]][key] = value
  730. elif self._tables:
  731. if not _is_table_or_aot(value): # if the value is a plain value
  732. for table in self._tables:
  733. # find the first table that allows plain values
  734. if any(not _is_table_or_aot(v) for _, v in table.items()):
  735. table[key] = value
  736. break
  737. else:
  738. self._tables[0][key] = value
  739. else:
  740. self._tables[0][key] = value
  741. else:
  742. self._container[key] = value
  743. self._internal_container[key] = value
  744. if key is not None:
  745. dict.__setitem__(self, key, value)
  746. def _remove_table(self, table: Table) -> None:
  747. """Remove table from the parent container"""
  748. self._tables.remove(table)
  749. for idx, item in enumerate(self._container._body):
  750. if item[1] is table:
  751. self._container._remove_at(idx)
  752. break
  753. def __delitem__(self, key: Key | str) -> None:
  754. if key not in self._tables_map:
  755. raise NonExistentKey(key)
  756. for i in reversed(self._tables_map[key]):
  757. table = self._tables[i]
  758. del table[key]
  759. if not table and len(self._tables) > 1:
  760. self._remove_table(table)
  761. del self._tables_map[key]
  762. del self._internal_container[key]
  763. if key is not None:
  764. dict.__delitem__(self, key)
  765. def __iter__(self) -> Iterator[str]:
  766. return iter(dict.keys(self))
  767. def __len__(self) -> int:
  768. return dict.__len__(self)
  769. def setdefault(self, key: Key | str, default: Any) -> Any:
  770. super().setdefault(key, default=default)
  771. return self[key]
  772. def ends_with_whitespace(it: Any) -> bool:
  773. """Returns ``True`` if the given item ``it`` is a ``Table`` or ``AoT`` object
  774. ending with a ``Whitespace``.
  775. """
  776. return (
  777. isinstance(it, Table) and isinstance(it.value._previous_item(), Whitespace)
  778. ) or (isinstance(it, AoT) and len(it) > 0 and isinstance(it[-1], Whitespace))