writeexts.py: convert 'mount' to context manager

2 files changed, 31 insertions, 40 deletions

diff --git a/openstack.write b/openstack.write
index 516fe367..b1941d3c 100755
--- a/openstack.write
+++ b/openstack.write
@@ -79,8 +79,7 @@ class OpenStackWriteExtension(morphlib.writeexts.WriteExtension):
     def set_extlinux_root_to_virtio(self, raw_disk):
         '''Re-configures extlinux to use virtio disks'''
         self.status(msg='Updating extlinux.conf')
-        mp = self.mount(raw_disk)
-        try:
+        with self.mount(raw_disk) as mp:
             path = os.path.join(mp, 'extlinux.conf')
 
             with open(path) as f:
@@ -91,9 +90,6 @@ class OpenStackWriteExtension(morphlib.writeexts.WriteExtension):
             with open(path, "w") as f:
                 f.write(extlinux_conf)
 
-        finally:
-            self.unmount(mp)
-
     def get_openstack_parameters(self):
         '''Get the environment variables needed.
 
diff --git a/rawdisk.write b/rawdisk.write
index e1a75fe0..b17f8aa7 100755
--- a/rawdisk.write
+++ b/rawdisk.write
@@ -67,51 +67,46 @@ class RawDiskWriteExtension(morphlib.writeexts.WriteExtension):
     def upgrade_local_system(self, raw_disk, temp_root):
         self.complete_fstab_for_btrfs_layout(temp_root)
 
-        mp = self.mount(raw_disk)
+        with self.mount(raw_disk) as mp:
+            version_label = self.get_version_label(mp)
+            self.status(msg='Updating image to a new version with label %s' %
+                        version_label)
+
+            version_root = os.path.join(mp, 'systems', version_label)
+            os.mkdir(version_root)
+
+            old_orig = os.path.join(mp, 'systems', 'default', 'orig')
+            new_orig = os.path.join(version_root, 'orig')
+            cliapp.runcmd(
+                ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])
+
+            cliapp.runcmd(
+                ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
+                 temp_root + os.path.sep, new_orig])
+
+            self.create_run(version_root)
+
+            default_path = os.path.join(mp, 'systems', 'default')
+            if os.path.exists(default_path):
+                os.remove(default_path)
+            else:
+                # we are upgrading and old system that does
+                # not have an updated extlinux config file
+                if self.bootloader_config_is_wanted():
+                    self.generate_bootloader_config(mp)
+                self.install_bootloader(mp)
+            os.symlink(version_label, default_path)
 
-        version_label = self.get_version_label(mp)
-        self.status(msg='Updating image to a new version with label %s' %
-                    version_label)
-
-        version_root = os.path.join(mp, 'systems', version_label)
-        os.mkdir(version_root)
-
-        old_orig = os.path.join(mp, 'systems', 'default', 'orig')
-        new_orig = os.path.join(version_root, 'orig')
-        cliapp.runcmd(
-            ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])
-
-        cliapp.runcmd(
-            ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
-             temp_root + os.path.sep, new_orig])
-
-        self.create_run(version_root)
-
-        default_path = os.path.join(mp, 'systems', 'default')
-        if os.path.exists(default_path):
-            os.remove(default_path)
-        else:
-            # we are upgrading and old system that does
-            # not have an updated extlinux config file
             if self.bootloader_config_is_wanted():
-                self.generate_bootloader_config(mp)
-            self.install_bootloader(mp)
-        os.symlink(version_label, default_path)
-
-        if self.bootloader_config_is_wanted():
-            self.install_kernel(version_root, temp_root)
-
-        self.unmount(mp)
+                self.install_kernel(version_root, temp_root)
 
     def get_version_label(self, mp):
         version_label = os.environ.get('VERSION_LABEL')
 
         if version_label is None:
-            self.unmount(mp)
             raise cliapp.AppException('VERSION_LABEL was not given')
 
         if os.path.exists(os.path.join(mp, 'systems', version_label)):
-            self.unmount(mp)
             raise cliapp.AppException('VERSION_LABEL %s already exists'
                                       % version_label)