From ec69fe9031b3f855a61838c60e284f469dc3ea63 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Randy=20D=C3=B6ring?=
 <30527984+radoering@users.noreply.github.com>
Date: Sat, 18 Jun 2022 14:58:47 +0200
Subject: [PATCH] Keep consistent line endings when changing files and use
 os.linesep as default for new files

---
 tests/test_toml_file.py | 43 +++++++++++++++++++++++++++++++++++++++++
 tomlkit/toml_file.py    | 29 +++++++++++++++++++++++++--
 2 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/tests/test_toml_file.py b/tests/test_toml_file.py
index e5b9b52..ae20656 100644
--- a/tests/test_toml_file.py
+++ b/tests/test_toml_file.py
@@ -62,3 +62,46 @@ def test_mixed_eol(tmpdir):
 
     with open(toml_path, "rb") as f:
         assert f.read() == b"a = 1\r\nrb = 2\n"
+
+
+def test_consistent_eol(tmpdir):
+    toml_path = str(tmpdir / "pyproject.toml")
+    with open(toml_path, "wb+") as f:
+        f.write(b"a = 1\r\nb = 2\r\n")
+
+    f = TOMLFile(toml_path)
+    content = f.read()
+    content["c"] = 3
+    f.write(content)
+
+    with open(toml_path, "rb") as f:
+        assert f.read() == b"a = 1\r\nb = 2\r\nc = 3\r\n"
+
+
+def test_consistent_eol_2(tmpdir):
+    toml_path = str(tmpdir / "pyproject.toml")
+    with open(toml_path, "wb+") as f:
+        f.write(b"a = 1\nb = 2\n")
+
+    f = TOMLFile(toml_path)
+    content = f.read()
+    content["c"] = 3
+    content["c"].trivia.trail = "\r\n"
+    f.write(content)
+
+    with open(toml_path, "rb") as f:
+        assert f.read() == b"a = 1\nb = 2\nc = 3\n"
+
+
+def test_default_eol_is_os_linesep(tmpdir):
+    toml_path = str(tmpdir / "pyproject.toml")
+    f = TOMLFile(toml_path)
+    content = TOMLDocument()
+    content.append("a", 1)
+    content["a"].trivia.trail = "\n"
+    content.append("b", 2)
+    content["b"].trivia.trail = "\r\n"
+    f.write(content)
+    linesep = os.linesep.encode()
+    with open(toml_path, "rb") as f:
+        assert f.read() == b"a = 1" + linesep + b"b = 2" + linesep
diff --git a/tomlkit/toml_file.py b/tomlkit/toml_file.py
index 47c9f2c..5b28cb0 100644
--- a/tomlkit/toml_file.py
+++ b/tomlkit/toml_file.py
@@ -1,3 +1,6 @@
+import os
+import re
+
 from .api import loads
 from .toml_document import TOMLDocument
 
@@ -11,13 +14,35 @@ class TOMLFile:
 
     def __init__(self, path: str) -> None:
         self._path = path
+        self._linesep = os.linesep
 
     def read(self) -> TOMLDocument:
         """Read the file content as a :class:`tomlkit.toml_document.TOMLDocument`."""
         with open(self._path, encoding="utf-8", newline="") as f:
-            return loads(f.read())
+            content = f.read()
+
+            # check if consistent line endings
+            num_newline = content.count("\n")
+            if num_newline > 0:
+                num_win_eol = content.count("\r\n")
+                if num_win_eol == num_newline:
+                    self._linesep = "\r\n"
+                elif num_win_eol == 0:
+                    self._linesep = "\n"
+                else:
+                    self._linesep = "mixed"
+
+            return loads(content)
 
     def write(self, data: TOMLDocument) -> None:
         """Write the TOMLDocument to the file."""
+        content = data.as_string()
+
+        # apply linesep
+        if self._linesep == "\n":
+            content = content.replace("\r\n", "\n")
+        elif self._linesep == "\r\n":
+            content = re.sub(r"(?<!\r)\n", "\r\n", content)
+
         with open(self._path, "w", encoding="utf-8", newline="") as f:
-            f.write(data.as_string())
+            f.write(content)