From 6c8a58ba84f74a26e8706fb9cf501a5f9ff57a3d Mon Sep 17 00:00:00 2001 From: Fabio Caccamo Date: Mon, 11 Dec 2023 12:17:38 +0100 Subject: [PATCH] Add possibility to write files atomically (`fsutil.write_file(path, content, atomic=True)`). #91 --- README.md | 4 ++-- fsutil/__init__.py | 58 ++++++++++++++++++++++++++++++++++++++++++---- tests/test.py | 46 ++++++++++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index dbe4dd8..67eed65 100644 --- a/README.md +++ b/README.md @@ -714,14 +714,14 @@ path_names = fsutil.split_path(path) ```python # Write file with the specified content at the given path. -fsutil.write_file(path, content, append=False, encoding="utf-8") +fsutil.write_file(path, content, append=False, encoding="utf-8", atomic=False) ``` #### `write_file_json` ```python # Write a json file at the given path with the specified data encoded in json format. -fsutil.write_file_json(path, data, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, cls=None, indent=None, separators=None, default=None, sort_keys=False) +fsutil.write_file_json(path, data, encoding="utf-8", atomic=False, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, cls=None, indent=None, separators=None, default=None, sort_keys=False) ``` ## Testing diff --git a/fsutil/__init__.py b/fsutil/__init__.py index 4a24fac..89c1e10 100644 --- a/fsutil/__init__.py +++ b/fsutil/__init__.py @@ -1314,12 +1314,50 @@ def split_path(path: PathIn) -> list[str]: return names +def _write_file_atomic( + path: PathIn, + content: str, + *, + append: bool = False, + encoding: str = "utf-8", +) -> None: + mode = "a" if append else "w" + if append: + content = read_file(path, encoding=encoding) + content + dirpath, _ = split_filepath(path) + with tempfile.NamedTemporaryFile( + mode=mode, + dir=dirpath, + delete=False, + encoding=encoding, + ) as file: + file.write(content) + file.flush() + os.fsync(file.fileno()) + temp_path = file.name + os.replace(temp_path, path) + remove_file(temp_path) + + +def _write_file_non_atomic( + path: PathIn, + content: str, + *, + append: bool = False, + encoding: str = "utf-8", +): + mode = "a" if append else "w" + with open(path, mode, encoding=encoding) as file: + file.write(content) + + def write_file( path: PathIn, content: str, *, append: bool = False, encoding: str = "utf-8", + atomic: bool = False, ) -> None: """ Write file with the specified content at the given path. @@ -1327,14 +1365,20 @@ def write_file( path = _get_path(path) assert_not_dir(path) make_dirs_for_file(path) - mode = "a" if append else "w" - with open(path, mode, encoding=encoding) as file: - file.write(content) + write_file_func = _write_file_atomic if atomic else _write_file_non_atomic + write_file_func( + path, + content, + append=append, + encoding=encoding, + ) def write_file_json( path: PathIn, data: Any, + encoding: str = "utf-8", + atomic: bool = False, **kwargs: Any, ) -> None: """ @@ -1351,4 +1395,10 @@ def default_encoder(obj: Any) -> Any: kwargs.setdefault("default", default_encoder) content = json.dumps(data, **kwargs) - write_file(path, content) + write_file( + path, + content, + append=False, + encoding=encoding, + atomic=atomic, + ) diff --git a/tests/test.py b/tests/test.py index f62d7a2..17afbf5 100644 --- a/tests/test.py +++ b/tests/test.py @@ -1214,6 +1214,17 @@ def test_write_file(self): fsutil.write_file(self.temp_path("a/b/c.txt"), content="Hello Jupiter") self.assertEqual(fsutil.read_file(path), "Hello Jupiter") + def test_write_file_atomic(self): + path = self.temp_path("a/b/c.txt") + fsutil.write_file( + self.temp_path("a/b/c.txt"), content="Hello World", atomic=True + ) + self.assertEqual(fsutil.read_file(path), "Hello World") + fsutil.write_file( + self.temp_path("a/b/c.txt"), content="Hello Jupiter", atomic=True + ) + self.assertEqual(fsutil.read_file(path), "Hello Jupiter") + def test_write_file_with_filename_only(self): path = "document.txt" fsutil.write_file(path, content="Hello World") @@ -1242,6 +1253,27 @@ def test_write_file_json(self): ), ) + def test_write_file_json_atomic(self): + path = self.temp_path("a/b/c.json") + now = datetime.now() + dec = Decimal("3.33") + data = { + "test": "Hello World", + "test_datetime": now, + "test_decimal": dec, + } + fsutil.write_file_json(self.temp_path("a/b/c.json"), data=data, atomic=True) + self.assertEqual( + fsutil.read_file(path), + ( + "{" + f'"test": "Hello World", ' + f'"test_datetime": "{now.isoformat()}", ' + f'"test_decimal": "{dec}"' + "}" + ), + ) + def test_write_file_with_append(self): path = self.temp_path("a/b/c.txt") fsutil.write_file(self.temp_path("a/b/c.txt"), content="Hello World") @@ -1251,6 +1283,20 @@ def test_write_file_with_append(self): ) self.assertEqual(fsutil.read_file(path), "Hello World - Hello Sun") + def test_write_file_with_append_atomic(self): + path = self.temp_path("a/b/c.txt") + fsutil.write_file( + self.temp_path("a/b/c.txt"), content="Hello World", atomic=True + ) + self.assertEqual(fsutil.read_file(path), "Hello World") + fsutil.write_file( + self.temp_path("a/b/c.txt"), + content=" - Hello Sun", + append=True, + atomic=True, + ) + self.assertEqual(fsutil.read_file(path), "Hello World - Hello Sun") + if __name__ == "__main__": unittest.main()