From e9ff02375f97cea418c85f10f7fee88e2377ad68 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Thu, 30 Nov 2017 21:52:51 -0800
Subject: [PATCH] More letter updates

---
 Scripts/functions/disk.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py
index f19bb81a..4a265300 100644
--- a/Scripts/functions/disk.py
+++ b/Scripts/functions/disk.py
@@ -4,13 +4,16 @@ from functions.common import *
 import partition_uids
 
 def assign_volume_letters():
+    with open(DISKPART_SCRIPT, 'w') as script:
+        for vol in get_volumes():
+            script.write('select volume {Number}\n'.format(**vol))
+            script.write('assign\n')
+    
+    # Remove current letters
     remove_volume_letters()
+    
+    # Run script
     try:
-        # Run script
-        with open(DISKPART_SCRIPT, 'w') as script:
-            for vol in get_volumes():
-                script.write('select volume {Number}\n'.format(**vol))
-                script.write('assign\n')
         run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
     except subprocess.CalledProcessError:
         pass
@@ -303,13 +306,14 @@ def reassign_volume_letter(letter, new_letter='I'):
 def remove_volume_letters(keep=None):
     if not keep:
         keep = ''
+    with open(DISKPART_SCRIPT, 'w') as script:
+        for vol in get_volumes():
+            if vol['Letter'].upper() != keep.upper():
+                script.write('select volume {Number}\n'.format(**vol))
+                script.write('remove noerr\n')
+    
+    # Run script
     try:
-        # Run script
-        with open(DISKPART_SCRIPT, 'w') as script:
-            for vol in get_volumes():
-                if vol['Letter'].upper() != keep.upper():
-                    script.write('select volume {Number}\n'.format(**vol))
-                    script.write('remove noerr\n')
         run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
     except subprocess.CalledProcessError:
         pass