From ea41eefa529dc0e5bbbfddb75e816e72e0daf0b0 Mon Sep 17 00:00:00 2001
From: Trim21 <i@trim21.me>
Date: Wed, 30 Sep 2020 02:17:00 +0800
Subject: [PATCH] preserve line breaking in TOMLFile

---
 tests/test_toml_file.py | 27 +++++++++++++++++++++++++++
 tomlkit/toml_file.py    | 15 +++++++++------
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/tests/test_toml_file.py b/tests/test_toml_file.py
index c7959654..8fa2f8a9 100644
--- a/tests/test_toml_file.py
+++ b/tests/test_toml_file.py
@@ -23,3 +23,30 @@ def test_toml_file(example):
     finally:
         with io.open(toml_file, "w", encoding="utf-8") as f:
             assert f.write(original_content)
+
+
+def test_keep_old_eol(tmpdir):
+    toml_path = str(tmpdir / "pyproject.toml")
+    with io.open(toml_path, "wb+") as f:
+        f.write(b"a = 1\r\nb = 2\r\n")
+
+    f = TOMLFile(toml_path)
+    content = f.read()
+    assert f._line_ending == "\r\n"
+    content["b"] = 3
+    f.write(content)
+
+    with io.open(toml_path, "rb") as f:
+        assert f.read() == b"a = 1\r\nb = 3\r\n"
+
+
+def test_mixed_eol(tmpdir):
+    toml_path = str(tmpdir / "pyproject.toml")
+    with io.open(toml_path, "wb+") as f:
+        f.write(b"a = 1\rb = 2\r\nc = 3\n")
+
+    f = TOMLFile(toml_path)
+    f.write(f.read())
+
+    with io.open(toml_path, "rb") as f:
+        assert f.read() == "a = 1{0}b = 2{0}c = 3{0}".format(os.linesep).encode()
diff --git a/tomlkit/toml_file.py b/tomlkit/toml_file.py
index 3b416664..809948c7 100644
--- a/tomlkit/toml_file.py
+++ b/tomlkit/toml_file.py
@@ -1,7 +1,5 @@
 import io
-
-from typing import Any
-from typing import Dict
+import os
 
 from .api import loads
 from .toml_document import TOMLDocument
@@ -14,11 +12,16 @@ class TOMLFile(object):
 
     def __init__(self, path):  # type: (str) -> None
         self._path = path
+        self._line_ending = os.linesep
 
     def read(self):  # type: () -> TOMLDocument
-        with io.open(self._path, encoding="utf-8") as f:
-            return loads(f.read())
+        with io.open(self._path, "r", encoding="utf-8", newline=None) as f:
+            content = f.read()
+            if not isinstance(f.newlines, tuple):
+                # mixed line endings
+                self._line_ending = f.newlines
+            return loads(content)
 
     def write(self, data):  # type: (TOMLDocument) -> None
-        with io.open(self._path, "w", encoding="utf-8") as f:
+        with io.open(self._path, "w", encoding="utf-8", newline=self._line_ending) as f:
             f.write(data.as_string())