diff --git a/tests/test_atomic_file.py b/tests/test_atomic_file.py index 010f899034..fbf522be95 100644 --- a/tests/test_atomic_file.py +++ b/tests/test_atomic_file.py @@ -3,9 +3,13 @@ from pathlib import Path import pytest +import typer +from typer.testing import CliRunner from . import atomic_write_example as mod +runner = CliRunner() + def test_atomic_write(tmp_path: Path) -> None: original_content = "existing-content\n" @@ -89,6 +93,41 @@ def test_atomic_api(tmp_path: Path) -> None: assert output_file.read_text(encoding="utf-8") == "atomic-api-done\n" +@pytest.mark.parametrize("lazy", [True, False], ids=["lazy", "eager"]) +@pytest.mark.parametrize( + "destination_exists", [True, False], ids=["existing", "missing"] +) +def test_atomic_write_callback_failure( + tmp_path: Path, lazy: bool, destination_exists: bool +) -> None: + original_content = "existing-content\n" + output_file = tmp_path / "atomic-failure-target.txt" + if destination_exists: + output_file.write_text(original_content, encoding="utf-8") + initial_entries = set(tmp_path.iterdir()) + + app = typer.Typer() + + @app.command() + def write_atomic_failure( + config: typer.FileTextWrite = typer.Option(..., atomic=True, lazy=lazy), + ) -> None: + config.write("partial-content\n") + config.flush() + raise RuntimeError("callback failed") + + result = runner.invoke(app, [f"--config={output_file}"]) + + assert result.exit_code == 1 + assert isinstance(result.exception, RuntimeError) + assert str(result.exception) == "callback failed" + if destination_exists: + assert output_file.read_text(encoding="utf-8") == original_content + else: + assert not output_file.exists() + assert set(tmp_path.iterdir()) == initial_entries + + @pytest.mark.parametrize( ("command_name", "expected_message"), [ diff --git a/typer/_click/_compat.py b/typer/_click/_compat.py index 6ed0ceb8ac..491279afe3 100644 --- a/typer/_click/_compat.py +++ b/typer/_click/_compat.py @@ -416,7 +416,13 @@ def close(self, delete: bool = False) -> None: if self.closed: return # pragma: no cover self._f.close() - os.replace(self._tmp_filename, self._real_filename) + if delete: + try: + os.unlink(self._tmp_filename) + except OSError: # pragma: no cover + pass + else: + os.replace(self._tmp_filename, self._real_filename) self.closed = True def __getattr__(self, name: str) -> Any: diff --git a/typer/_click/types.py b/typer/_click/types.py index f49244c6c5..6fad40a017 100644 --- a/typer/_click/types.py +++ b/typer/_click/types.py @@ -539,7 +539,10 @@ def convert( ) if ctx is not None: - ctx.call_on_close(lf.close_intelligently) + if self.atomic: + ctx.with_resource(lf) + else: + ctx.call_on_close(lf.close_intelligently) return cast("IO[Any]", lf) @@ -554,7 +557,10 @@ def convert( # type is used with prompts. if ctx is not None: if should_close: - ctx.call_on_close(safecall(f.close)) + if self.atomic: + ctx.with_resource(f) + else: + ctx.call_on_close(safecall(f.close)) else: ctx.call_on_close(safecall(f.flush)) diff --git a/typer/_click/utils.py b/typer/_click/utils.py index ac8e5ba3f2..f310672de5 100644 --- a/typer/_click/utils.py +++ b/typer/_click/utils.py @@ -182,7 +182,10 @@ def __exit__( exc_value: BaseException | None, tb: TracebackType | None, ) -> None: - self.close_intelligently() + if self.atomic and self.should_close and self._f is not None: + self._f.__exit__(exc_type, exc_value, tb) + else: + self.close_intelligently() def __iter__(self) -> Iterator[AnyStr]: self.open()