From 7151739ad31928ff0f4206a13dcddf00919f705f Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Wed, 22 Feb 2017 05:42:40 -0800
Subject: [PATCH] Added transaction support to write & delete files. * if error
 occurs, disk image returned to state it was before the call

---
 atrcopy/diskimages.py | 68 ++++++++++++++++++++++++++++---------------
 1 file changed, 45 insertions(+), 23 deletions(-)

diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py
index 61a8199..ba33230 100644
--- a/atrcopy/diskimages.py
+++ b/atrcopy/diskimages.py
@@ -377,6 +377,14 @@ class DiskImageBase(object):
 
     # file writing methods
 
+    def begin_transaction(self):
+        state = self.bytes[:], self.style[:]
+        return state
+
+    def rollback_transaction(self, state):
+        self.bytes[:], self.style[:] = state
+        return
+
     def write_file(self, filename, filetype, data):
         """Write data to a file on disk
 
@@ -384,18 +392,25 @@ class DiskImageBase(object):
         not enough space on disk or a free entry is not available in the
         catalog.
         """
-        directory = self.directory_class(self.bytes_per_sector)
-        self.get_directory(directory)
-        dirent = directory.add_dirent(filename, filetype)
-        data = to_numpy(data)
-        sector_list = self.sector_list_class(self.bytes_per_sector, self.payload_bytes_per_sector, data, self.writeable_sector_class)
-        vtoc_segments = self.get_vtoc_segments()
-        vtoc = self.vtoc_class(self.bytes_per_sector, vtoc_segments)
-        directory.save_dirent(dirent, vtoc, sector_list)
-        self.write_sector_list(sector_list)
-        self.write_sector_list(vtoc)
-        self.write_sector_list(directory)
-        self.get_metadata()
+        state = self.begin_transaction()
+        try:
+            directory = self.directory_class(self.bytes_per_sector)
+            self.get_directory(directory)
+            dirent = directory.add_dirent(filename, filetype)
+            data = to_numpy(data)
+            sector_list = self.sector_list_class(self.bytes_per_sector, self.payload_bytes_per_sector, data, self.writeable_sector_class)
+            vtoc_segments = self.get_vtoc_segments()
+            vtoc = self.vtoc_class(self.bytes_per_sector, vtoc_segments)
+            directory.save_dirent(dirent, vtoc, sector_list)
+            self.write_sector_list(sector_list)
+            self.write_sector_list(vtoc)
+            self.write_sector_list(directory)
+            self.get_metadata()
+        except AtrError:
+            self.rollback_transaction(state)
+            raise
+        finally:
+            self.get_metadata()
 
     def write_sector_list(self, sector_list):
         for sector in sector_list:
@@ -404,17 +419,24 @@ class DiskImageBase(object):
             self.bytes[pos:pos + size] = sector.data
 
     def delete_file(self, filename):
-        directory = self.directory_class(self.bytes_per_sector)
-        self.get_directory(directory)
-        dirent = directory.find_dirent(filename)
-        sector_list = dirent.get_sector_list(self)
-        vtoc_segments = self.get_vtoc_segments()
-        vtoc = self.vtoc_class(self.bytes_per_sector, vtoc_segments)
-        directory.remove_dirent(dirent, vtoc, sector_list)
-        self.write_sector_list(sector_list)
-        self.write_sector_list(vtoc)
-        self.write_sector_list(directory)
-        self.get_metadata()
+        state = self.begin_transaction()
+        try:
+            directory = self.directory_class(self.bytes_per_sector)
+            self.get_directory(directory)
+            dirent = directory.find_dirent(filename)
+            sector_list = dirent.get_sector_list(self)
+            vtoc_segments = self.get_vtoc_segments()
+            vtoc = self.vtoc_class(self.bytes_per_sector, vtoc_segments)
+            directory.remove_dirent(dirent, vtoc, sector_list)
+            self.write_sector_list(sector_list)
+            self.write_sector_list(vtoc)
+            self.write_sector_list(directory)
+            self.get_metadata()
+        except AtrError:
+            self.rollback_transaction(state)
+            raise
+        finally:
+            self.get_metadata()
 
 
 class WriteableSector(object):