toml_file.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  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. content = content.replace("\r\n", "\n")
  30. elif num_win_eol == 0:
  31. self._linesep = "\n"
  32. else:
  33. self._linesep = "mixed"
  34. return loads(content)
  35. def write(self, data: TOMLDocument) -> None:
  36. """Write the TOMLDocument to the file."""
  37. content = data.as_string()
  38. # apply linesep
  39. if self._linesep == "\n":
  40. content = content.replace("\r\n", "\n")
  41. elif self._linesep == "\r\n":
  42. content = re.sub(r"(?<!\r)\n", "\r\n", content)
  43. with open(self._path, "w", encoding="utf-8", newline="") as f:
  44. f.write(content)