diff --git a/common/file_helpers.py b/common/file_helpers.py index a29eafdd9f..7524860f96 100644 --- a/common/file_helpers.py +++ b/common/file_helpers.py @@ -1,7 +1,8 @@ import os import shutil import tempfile -from atomicwrites import AtomicWriter +import contextlib +from typing import Optional def rm_not_exists_ok(path): @@ -71,19 +72,25 @@ def _get_fileobject_func(writer, temp_dir): return writer.get_fileobject(dir=temp_dir) return _get_fileobject -def atomic_write_on_fs_tmp(path, **kwargs): - """Creates an atomic writer using a temporary file in a temporary directory - on the same filesystem as path. - """ - # TODO(mgraczyk): This use of AtomicWriter relies on implementation details to set the temp - # directory. - writer = AtomicWriter(path, **kwargs) - return writer._open(_get_fileobject_func(writer, get_tmpdir_on_same_filesystem(path))) - - -def atomic_write_in_dir(path, **kwargs): - """Creates an atomic writer using a temporary file in the same directory - as the destination file. - """ - writer = AtomicWriter(path, **kwargs) - return writer._open(_get_fileobject_func(writer, os.path.dirname(path))) +@contextlib.contextmanager +def atomic_write_on_fs_tmp(path: str, mode: str = 'w', buffering: int = -1, encoding: Optional[str] = None, newline: Optional[str] = None): + """Write to a file atomically using a temporary file in a temporary directory on the same filesystem as path.""" + temp_dir = get_tmpdir_on_same_filesystem(path) + with tempfile.NamedTemporaryFile(mode=mode, buffering=buffering, encoding=encoding, newline=newline, dir=temp_dir, delete=False) as tmp_file: + yield tmp_file + tmp_file_name = tmp_file.name + os.replace(tmp_file_name, path) + +@contextlib.contextmanager +def atomic_write_in_dir(path: str, mode: str = 'w', buffering: int = -1, encoding: Optional[str] = None, newline: Optional[str] = None, + overwrite: bool = False): + """Write to a file atomically using a temporary file in the same directory as the destination file.""" + dir_name = os.path.dirname(path) + + if not overwrite and os.path.exists(path): + raise FileExistsError(f"File '{path}' already exists. To overwrite it, set 'overwrite' to True.") + + with tempfile.NamedTemporaryFile(mode=mode, buffering=buffering, encoding=encoding, newline=newline, dir=dir_name, delete=False) as tmp_file: + yield tmp_file + tmp_file_name = tmp_file.name + os.replace(tmp_file_name, path) diff --git a/poetry.lock b/poetry.lock index b58c8c9aa9..8443506e4c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:411564561a9b238f4c2e64f66a50a112aa0a5372e559e9e187979a87e1d026f4 -size 435710 +oid sha256:8fe4d2310328e34ec37f396db589ef55ad34a3d27c7add141abd9235fd38eb11 +size 434974 diff --git a/pyproject.toml b/pyproject.toml index 336c440f6b..98091e0088 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,7 +108,6 @@ websocket_client = "*" polyline = "*" # these should be removed -atomicwrites = "*" markdown-it-py = "*" timezonefinder = "*" pycurl = "*" @@ -154,7 +153,6 @@ sphinx-rtd-theme = "*" sphinx-sitemap = "*" tabulate = "*" tenacity = "*" -types-atomicwrites = "*" types-pycurl = "*" types-requests = "*" types-tabulate = "*"