toml_file.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import os
  2. import re
  3. from typing import TYPE_CHECKING
  4. from tomlkit.api import loads
  5. from tomlkit.toml_document import TOMLDocument
  6. if TYPE_CHECKING:
  7. from _typeshed import StrPath as _StrPath
  8. else:
  9. from typing import Union
  10. _StrPath = Union[str, os.PathLike]
  11. class TOMLFile:
  12. """
  13. Represents a TOML file.
  14. :param path: path to the TOML file
  15. """
  16. def __init__(self, path: _StrPath) -> None:
  17. self._path = path
  18. self._linesep = os.linesep
  19. def read(self) -> TOMLDocument:
  20. """Read the file content as a :class:`tomlkit.toml_document.TOMLDocument`."""
  21. with open(self._path, encoding="utf-8", newline="") as f:
  22. content = f.read()
  23. # check if consistent line endings
  24. num_newline = content.count("\n")
  25. if num_newline > 0:
  26. num_win_eol = content.count("\r\n")
  27. if num_win_eol == num_newline:
  28. self._linesep = "\r\n"
  29. elif num_win_eol == 0:
  30. self._linesep = "\n"
  31. else:
  32. self._linesep = "mixed"
  33. return loads(content)
  34. def write(self, data: TOMLDocument) -> None:
  35. """Write the TOMLDocument to the file."""
  36. content = data.as_string()
  37. # apply linesep
  38. if self._linesep == "\n":
  39. content = content.replace("\r\n", "\n")
  40. elif self._linesep == "\r\n":
  41. content = re.sub(r"(?<!\r)\n", "\r\n", content)
  42. with open(self._path, "w", encoding="utf-8", newline="") as f:
  43. f.write(content)