Cut ~59s from every offline install by not shelling out to Get-Disk

ext4_symlink.py

@@ -191,14 +191,57 @@ def _ps_single_quote(s: str) -> str:
     return s.replace("'", "''")
 
 
-def _disk_number(vhd_path: str) -> int | None:
-    """OS disk number of the attached VHD (== PhysicalDriveN index)."""
+_ASSOC_DISK_RE = re.compile(r"Associated disk#\s*:\s*(\d+)", re.IGNORECASE)
+_ASSOC_DISK_FIELD_RE = re.compile(r"Associated disk#", re.IGNORECASE)
+
+
+def _disk_number_via_diskpart(vhd_path: str) -> tuple[int | None, bool]:
+    """``(disk_number, parsed)`` from ``diskpart``'s own ``detail vdisk``.
+
+    diskpart already knows which disk it just attached and prints it as
+    ``Associated disk#: 2``, so asking it costs one diskpart run (~1.6s) where
+    the ``Get-Disk`` route costs ~17s -- almost all of which is PowerShell cold-
+    loading the Storage/CIM providers rather than any disk work. That fires
+    twice per attach/detach cycle and twice per disk per install, so it was the
+    single largest component of the user's wait.
+
+    ``parsed`` says whether the output was understood at all, which is what lets
+    the caller fall back safely: the field name is English diskpart text, so a
+    localised Windows must not silently look like "no disk attached".
+    """
+    r = _diskpart('select vdisk file="%s"\ndetail vdisk\n' % vhd_path)
+    out = r.stdout or ""
+    m = _ASSOC_DISK_RE.search(out)
+    if m:
+        return int(m.group(1)), True
+    if _ASSOC_DISK_FIELD_RE.search(out):
+        return None, True          # field present, no number => detached
+    return None, False             # unrecognised output => caller falls back
+
+
+def _disk_number_via_get_disk(vhd_path: str) -> int | None:
+    """OS disk number via PowerShell ``Get-Disk``. Slow (~17s) but locale-proof."""
     r = _run(["powershell", "-NoProfile", "-Command",
               "(Get-Disk | Where-Object { $_.Location -eq '%s' }).Number" % _ps_single_quote(vhd_path)])
     out = (r.stdout or "").strip()
     return int(out) if out.isdigit() else None
 
 
+def _disk_number(vhd_path: str) -> int | None:
+    """OS disk number of the attached VHD (== PhysicalDriveN index), or None.
+
+    Fast path first, ``Get-Disk`` only when diskpart's output isn't recognised.
+    Both agree in both states -- attached and detached -- which matters because
+    :func:`_detach` treats ``None`` as proof the disk is gone.
+    """
+    number, parsed = _disk_number_via_diskpart(vhd_path)
+    if parsed:
+        return number
+    logger.debug("diskpart 'detail vdisk' not recognised (localised Windows?); "
+                 "falling back to Get-Disk for %s", vhd_path)
+    return _disk_number_via_get_disk(vhd_path)
+
+
 def _cyg_device(disk_number: int, offset: int) -> str:
     # Cygwin maps PhysicalDriveN -> /dev/sd<a+N>; ?offset= selects the ext4
     # partition without relying on Cygwin's own partition detection.

tests/test_disk_number.py

@@ -0,0 +1,80 @@
+"""_disk_number: fast diskpart parse, with a Get-Disk fallback for localisation.
+
+Pins the behaviour that made the fast path safe to adopt -- it must agree with
+Get-Disk in BOTH states, because _detach() treats None as proof the disk is gone.
+A fast path that returned None for an *attached* disk would make _detach report
+success while the image stayed mounted.
+"""
+from unittest import mock
+
+import ext4_symlink as es
+
+
+ATTACHED = """
+DiskPart successfully selected the virtual disk file.
+
+Device type ID: 3 (Unknown)
+Virtual size:  128 GB
+Physical size: 3530 MB
+Filename: C:\\ProgramData\\BlueStacks_nxt\\Engine\\Tiramisu64\\Data.vhdx
+Is Child: No
+Associated disk#: 2
+"""
+
+DETACHED = """
+DiskPart successfully selected the virtual disk file.
+
+Device type ID: 3 (Unknown)
+Virtual size:  128 GB
+Filename: C:\\ProgramData\\BlueStacks_nxt\\Engine\\Tiramisu64\\Data.vhdx
+Is Child: No
+Associated disk#:
+"""
+
+# A localised install prints the same data under translated field names.
+LOCALISED = """
+DiskPart hat die Datei für den virtuellen Datenträger ausgewählt.
+
+Datenträger 2
+"""
+
+
+def _completed(stdout):
+    return mock.Mock(stdout=stdout, stderr="", returncode=0)
+
+
+def test_attached_returns_disk_number_without_calling_get_disk():
+    with mock.patch.object(es, "_diskpart", return_value=_completed(ATTACHED)), \
+         mock.patch.object(es, "_disk_number_via_get_disk") as slow:
+        assert es._disk_number("X.vhdx") == 2
+    slow.assert_not_called()
+
+
+def test_detached_returns_none_without_calling_get_disk():
+    # _detach() relies on this: None must mean "gone", not "couldn't tell".
+    with mock.patch.object(es, "_diskpart", return_value=_completed(DETACHED)), \
+         mock.patch.object(es, "_disk_number_via_get_disk") as slow:
+        assert es._disk_number("X.vhdx") is None
+    slow.assert_not_called()
+
+
+def test_unrecognised_output_falls_back_to_get_disk():
+    # Must NOT be mistaken for "detached" -- that would let _detach() claim
+    # success on a still-mounted image on any non-English Windows.
+    with mock.patch.object(es, "_diskpart", return_value=_completed(LOCALISED)), \
+         mock.patch.object(es, "_disk_number_via_get_disk", return_value=7) as slow:
+        assert es._disk_number("X.vhdx") == 7
+    slow.assert_called_once()
+
+
+def test_parse_flag_distinguishes_detached_from_unparseable():
+    with mock.patch.object(es, "_diskpart", return_value=_completed(DETACHED)):
+        assert es._disk_number_via_diskpart("X.vhdx") == (None, True)
+    with mock.patch.object(es, "_diskpart", return_value=_completed(LOCALISED)):
+        assert es._disk_number_via_diskpart("X.vhdx") == (None, False)
+
+
+def test_case_insensitive_field_match():
+    with mock.patch.object(es, "_diskpart",
+                           return_value=_completed("associated DISK#:  11\n")):
+        assert es._disk_number_via_diskpart("X.vhdx") == (11, True)