diff --git a/scone/default/utensils/basic_utensils.py b/scone/default/utensils/basic_utensils.py index 6cbb239..09921bb 100644 --- a/scone/default/utensils/basic_utensils.py +++ b/scone/default/utensils/basic_utensils.py @@ -35,10 +35,13 @@ from scone.sous.utensils import Utensil, Worktop class WriteFile(Utensil): path: str mode: int + atomic: bool = attr.ib(default=True) async def execute(self, channel: Channel, worktop): oldumask = os.umask(0) - fdnum = os.open(self.path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, self.mode) + temp_path = self.path + "._scone-part" + write_path = temp_path if self.atomic else self.path + fdnum = os.open(write_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, self.mode) os.umask(oldumask) with open(fdnum, "wb") as file: @@ -49,6 +52,9 @@ class WriteFile(Utensil): assert isinstance(next_chunk, bytes) file.write(next_chunk) + if self.atomic: + shutil.move(temp_path, self.path) + await channel.send("OK")