From dbe4a342cc00baaf3b3da915034e5c7dc3b0a018 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 26 Aug 2023 14:30:22 -0700
Subject: [PATCH] Fix source_parts usage

Addresses issue #221
---
 scripts/wk/clone/block_pair.py | 18 +++++++++++++-----
 scripts/wk/clone/state.py      |  4 ++--
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/scripts/wk/clone/block_pair.py b/scripts/wk/clone/block_pair.py
index 5e330a51..d9d5b2b5 100644
--- a/scripts/wk/clone/block_pair.py
+++ b/scripts/wk/clone/block_pair.py
@@ -244,7 +244,7 @@ class BlockPair():
 
 
 # Functions
-def add_clone_block_pairs(state) -> None:
+def add_clone_block_pairs(state) -> list[hw_disk.Disk]:
   """Add device to device block pairs and set settings if necessary."""
   source_sep = get_partition_separator(state.source.path.name)
   dest_sep = get_partition_separator(state.destination.path.name)
@@ -255,6 +255,7 @@ def add_clone_block_pairs(state) -> None:
 
   # Add pairs from previous run
   if settings['Partition Mapping']:
+    source_parts = []
     for part_map in settings['Partition Mapping']:
       bp_source = hw_disk.Disk(
         f'{state.source.path}{source_sep}{part_map[0]}',
@@ -262,8 +263,9 @@ def add_clone_block_pairs(state) -> None:
       bp_dest = pathlib.Path(
         f'{state.destination.path}{dest_sep}{part_map[1]}',
         )
+      source_parts.append(bp_source)
       state.add_block_pair(bp_source, bp_dest)
-      return
+    return source_parts
 
   # Add pairs from selection
   source_parts = menus.select_disk_parts('Clone', state.source)
@@ -271,7 +273,7 @@ def add_clone_block_pairs(state) -> None:
     # Whole disk (or single partition via args), skip settings
     bp_dest = state.destination.path
     state.add_block_pair(state.source, bp_dest)
-    return
+    return source_parts
 
   # New run, use new settings file
   settings['Needs Format'] = True
@@ -306,13 +308,19 @@ def add_clone_block_pairs(state) -> None:
   # Save settings
   state.save_settings(settings)
 
+  # Done
+  return source_parts
 
-def add_image_block_pairs(state) -> None:
+
+def add_image_block_pairs(state) -> list[hw_disk.Disk]:
   """Add device to image file block pairs."""
   source_parts = menus.select_disk_parts(state.mode, state.source)
   for part in source_parts:
     state.add_block_pair(part, state.destination)
 
+  # Done
+  return source_parts
+
 
 def build_block_pair_report(block_pairs, settings) -> list:
   """Build block pair report, returns list."""
@@ -378,7 +386,7 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details) -> str:
 
   # Set dest type
   if re.match(r'^0x\w+$', source_type):
-    # Both source and dest are MBR
+    # Source is a MBR type
     source_table_type = 'MBR'
     if table_type == 'MBR':
       dest_type = source_type.replace('0x', '').lower()
diff --git a/scripts/wk/clone/state.py b/scripts/wk/clone/state.py
index bb0d6c26..f0dd9c66 100644
--- a/scripts/wk/clone/state.py
+++ b/scripts/wk/clone/state.py
@@ -378,9 +378,9 @@ class State():
 
     # Add block pairs
     if self.mode == 'Clone':
-      add_clone_block_pairs(self)
+      source_parts = add_clone_block_pairs(self)
     else:
-      add_image_block_pairs(self)
+      source_parts = add_image_block_pairs(self)
 
     # Update SMART data
     ## TODO: Verify if needed