From 29d4e80f7efed8fe3832fa4adf5128365636b0cc Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Fri, 20 May 2022 19:44:12 -0700
Subject: [PATCH] Refactor get_disks_macos() to improve performance

Found a way to easily exclude disk images from the list under El Cap
---
 scripts/wk/hw/disk.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py
index 45e5c50b..d58d5734 100644
--- a/scripts/wk/hw/disk.py
+++ b/scripts/wk/hw/disk.py
@@ -3,6 +3,7 @@
 
 import logging
 import pathlib
+import platform
 import plistlib
 import re
 
@@ -331,12 +332,12 @@ def get_disks_macos() -> list[Disk]:
   cmd = ['diskutil', 'list', '-plist', 'physical']
   disks = []
 
+  # El Capitan workaround
+  if platform.mac_ver()[0].startswith('10.11'):
+    cmd.pop()
+
   # Get info from diskutil
   proc = run_program(cmd, encoding=None, errors=None, check=False)
-  if proc.returncode != 0:
-    # Assuming we're running on an older macOS version
-    cmd.pop(-1)
-    proc = run_program(cmd, encoding=None, errors=None, check=False)
 
   # Parse plist data
   try:
@@ -347,13 +348,20 @@ def get_disks_macos() -> list[Disk]:
     return disks
 
   # Add valid disks
-  for disk in plist_data['WholeDisks']:
-    disks.append(Disk(f'/dev/{disk}'))
+  for disk in plist_data['AllDisksAndPartitions']:
+    name = disk['DeviceIdentifier']
+    if name not in plist_data['WholeDisks']:
+      # Only check "WholeDisks"
+      continue
+    if not disk['Content']:
+      # This lists GPT or MBR for device, blank should mean it's an image
+      continue
+    disks.append(Disk(f'/dev/{name}'))
 
   # Remove virtual disks
   # TODO: Test more to figure out why some drives are being marked 'Unknown'
   disks = [
-    d for d in disks if d.details.get('VirtualOrPhysical') != 'Virtual'
+    d for d in disks if d.raw_details.get('VirtualOrPhysical') != 'Virtual'
     ]
 
   # Done