Merge pull request #122 from timlnx/add-windows-support

Add Windows device capacity support via DeviceIoControl, closes #52

Author: timlnx

.github/workflows/python.yml

@@ -11,7 +11,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
-        os: ["macos-latest", "ubuntu-latest"]
+        os: ["macos-latest", "ubuntu-latest", "windows-latest"]
     runs-on: ${{ matrix.os }}
     steps:
       - name: "GitHub Checks it out :sunglasses-face:"
@@ -24,6 +24,7 @@ jobs:
           cache: 'pip'
 
       - name: Verify Test Case Names Unique
+        if: runner.os != 'Windows'
         run: |
           ./tests/test_unique_testcase_names.sh
 
@@ -34,8 +35,8 @@ jobs:
 
       - name: Pre-Tests code smell validation
         run: |
-          pycodestyle -v --ignore=E501,E722 bitmath/__init__.py tests/*.py
-          flake8 --select=F bitmath/__init__.py tests/*.py
+          pycodestyle -v --ignore=E501,E722 bitmath/__init__.py tests/
+          flake8 --select=F bitmath/__init__.py tests/
 
       - name: Run Unit Tests
         run: |

bitmath/__init__.py

@@ -58,13 +58,19 @@
 from collections.abc import Generator, Iterable, Iterator
 from typing import IO, Any
 
-# For device capacity reading in query_device_capacity(). Only supported
-# on posix systems for now. Will be addressed in issue #52 on GitHub.
+# For device capacity reading in query_device_capacity().
 if os.name == 'posix':
     import stat
     import fcntl
     import struct
+elif os.name == 'nt':
+    import ctypes
+    import ctypes.wintypes
+    import msvcrt
 
+#: Platforms where :func:`query_device_capacity` is supported.
+#: Corresponds to possible values of :data:`os.name`.
+SUPPORTED_PLATFORMS = frozenset({'posix', 'nt'})
 
 __all__ = ['Bit', 'Byte', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB',
            'kB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB', 'Kib',
@@ -1285,6 +1291,56 @@ def best_prefix(bytes: Bitmath | int | float, system: int = NIST) -> Bitmath:
     return Byte(value).best_prefix(system=system)
 
 
+def _query_device_capacity_windows(device_fd: IO[Any]) -> int:
+    """Return device capacity in bytes on Windows via DeviceIoControl.
+
+Windows physical disk paths look like ``\\\\.\\PhysicalDrive0``.
+Raises :class:`ValueError` if the file descriptor is not a physical device.
+Raises :class:`OSError` if the DeviceIoControl call fails.
+"""
+    if not device_fd.name.startswith('\\\\.\\'):
+        raise ValueError("The file descriptor provided is not of a device type")
+
+    IOCTL_DISK_GET_DRIVE_GEOMETRY_EX = 0x000700A0
+
+    class DISK_GEOMETRY(ctypes.Structure):
+        _fields_ = [
+            ('Cylinders', ctypes.c_longlong),
+            ('MediaType', ctypes.c_uint),
+            ('TracksPerCylinder', ctypes.c_ulong),
+            ('SectorsPerTrack', ctypes.c_ulong),
+            ('BytesPerSector', ctypes.c_ulong),
+        ]
+
+    class DISK_GEOMETRY_EX(ctypes.Structure):
+        _fields_ = [
+            ('Geometry', DISK_GEOMETRY),
+            ('DiskSize', ctypes.c_longlong),
+            ('Data', ctypes.c_byte * 1),
+        ]
+
+    geometry = DISK_GEOMETRY_EX()
+    bytes_returned = ctypes.wintypes.DWORD(0)
+    handle = msvcrt.get_osfhandle(device_fd.fileno())
+
+    result = ctypes.windll.kernel32.DeviceIoControl(
+        handle,
+        IOCTL_DISK_GET_DRIVE_GEOMETRY_EX,
+        None,
+        0,
+        ctypes.byref(geometry),
+        ctypes.sizeof(geometry),
+        ctypes.byref(bytes_returned),
+        None,
+    )
+
+    if not result:
+        error_code = ctypes.windll.kernel32.GetLastError()
+        raise OSError(f"DeviceIoControl failed with error code: {error_code}")
+
+    return geometry.DiskSize
+
+
 def query_device_capacity(device_fd: IO[Any]) -> Byte:
     """Create bitmath instances of the capacity of a system block device
 
@@ -1300,13 +1356,16 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte:
 * http://stackoverflow.com/a/9764508/263969
 
    :param file device_fd: A ``file`` object of the device to query the
-   capacity of (as in ``get_device_capacity(open("/dev/sda"))``).
+   capacity of. On Linux/macOS: ``open("/dev/sda", "rb")``. On Windows:
+   ``open(r'\\\\.\\PhysicalDrive0', 'rb')`` (requires administrator privileges).
 
    :return: a bitmath :class:`bitmath.Byte` instance equivalent to the
    capacity of the target device in bytes.
 """
-    if os.name != 'posix':
+    if os.name not in SUPPORTED_PLATFORMS:
         raise NotImplementedError(f"'bitmath.query_device_capacity' is not supported on this platform: {os.name}")
+    if os.name == 'nt':
+        return Byte(_query_device_capacity_windows(device_fd))
 
     s = os.stat(device_fd.name).st_mode
     if not stat.S_ISBLK(s):

tests/test_file_size.py

@@ -32,6 +32,7 @@
 from . import TestCase
 import bitmath
 import os
+import pathlib
 
 
 class TestFileSize(TestCase):
@@ -113,8 +114,8 @@ def test_listdir_nosymlinks(self):
 
         # Ensure the returned paths match the expected paths
         discovered_paths = [
-            contents[0][0],
-            contents[1][0],
+            pathlib.Path(contents[0][0]).as_posix(),
+            pathlib.Path(contents[1][0]).as_posix(),
         ]
         expected_paths = [
             'tests/listdir_nosymlinks/depth1/depth2/10_byte_file',
@@ -183,8 +184,8 @@ def test_listdir_symlinks_follow(self):
             'tests/listdir_symlinks/depth1/depth2/10_byte_file'
         ]
         discovered_paths = [
-            contents[0][0],
-            contents[1][0]
+            pathlib.Path(contents[0][0]).as_posix(),
+            pathlib.Path(contents[1][0]).as_posix()
         ]
         self.assertListEqualUnordered(discovered_paths, expected_paths)
 
@@ -243,7 +244,7 @@ def test_listdir_filtering_nosymlinks(self):
                                         filter='1024*'))
 
         # Ensure the returned path matches the expected path
-        self.assertEqual(contents[0][0],
+        self.assertEqual(pathlib.Path(contents[0][0]).as_posix(),
                          'tests/listdir_nosymlinks/depth1/depth2/1024_byte_file')
 
         # Ensure the measured size is what we expect

tests/test_query_device_capacity.py

@@ -31,7 +31,8 @@
 
 from . import TestCase
 import bitmath
-from unittest import mock
+import os
+from unittest import mock, skipUnless
 import struct
 from contextlib import ExitStack, contextmanager
 
@@ -50,8 +51,16 @@ def nested(*contexts):
 non_device_file = mock.MagicMock('file')
 non_device_file.name = "/home/"
 
+windows_device = mock.MagicMock('file')
+windows_device.fileno = mock.Mock(return_value=4)
+windows_device.name = r'\\.\PhysicalDrive0'
+
+windows_non_device = mock.MagicMock('file')
+windows_non_device.name = r'C:\somefile.txt'
+
 
 class TestQueryDeviceCapacity(TestCase):
+    @skipUnless(os.name == 'posix', 'fcntl is POSIX only')
     def test_query_device_capacity_linux_everything_is_wonderful(self):
         """query device capacity works on a happy Linux host"""
         with nested(
@@ -71,6 +80,7 @@ def test_query_device_capacity_linux_everything_is_wonderful(self):
             self.assertEqual(ioctl.call_count, 1)
             ioctl.assert_called_once_with(4, 0x80081272, struct.calcsize('L'))
 
+    @skipUnless(os.name == 'posix', 'fcntl is POSIX only')
     def test_query_device_capacity_mac_everything_is_wonderful(self):
         """query device capacity works on a happy Mac OS X host"""
         with nested(
@@ -100,6 +110,7 @@ def side_effect(*args, **kwargs):
             self.assertEqual(bytes, 1000000000000)
             self.assertEqual(ioctl.call_count, 2)
 
+    @skipUnless(os.name == 'posix', 'fcntl is POSIX only')
     def test_query_device_capacity_device_not_block(self):
         """query device capacity aborts if a non-block-device is provided"""
         with nested(
@@ -116,8 +127,29 @@ def test_query_device_capacity_device_not_block(self):
 
             self.assertEqual(ioctl.call_count, 0)
 
-    def test_query_device_capacity_non_posix_system_fails(self):
-        """query device capacity fails on a non-posix host"""
+    def test_query_device_capacity_windows_everything_is_wonderful(self):
+        """query device capacity works on a happy Windows host"""
+        expected_bytes = 1_000_000_000_000  # 1 TB
+
+        with mock.patch('bitmath._query_device_capacity_windows', return_value=expected_bytes):
+            with mock.patch('bitmath.os.name', 'nt'):
+                result = bitmath.query_device_capacity(windows_device)
+
+        self.assertEqual(result, bitmath.Byte(expected_bytes))
+
+    def test_query_device_capacity_windows_non_device_fails(self):
+        """query device capacity rejects a non-device path on Windows"""
         with mock.patch('bitmath.os.name', 'nt'):
+            with self.assertRaises(ValueError):
+                bitmath.query_device_capacity(windows_non_device)
+
+    def test_query_device_capacity_unsupported_platform_fails(self):
+        """query device capacity fails on an unsupported platform"""
+        # Derive a value that is guaranteed not to be in SUPPORTED_PLATFORMS.
+        unsupported = next(
+            p for p in ('os2', 'java', 'riscos', 'ce')
+            if p not in bitmath.SUPPORTED_PLATFORMS
+        )
+        with mock.patch('bitmath.os.name', unsupported):
             with self.assertRaises(NotImplementedError):
                 bitmath.query_device_capacity(device)