diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 154ea2f..acb9082 100644 --- a/.github/workflows/python.yml +++ b/.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: | diff --git a/bitmath/__init__.py b/bitmath/__init__.py index 08a0073..3739e1c 100644 --- a/bitmath/__init__.py +++ b/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): diff --git a/tests/test_file_size.py b/tests/test_file_size.py index ef0024d..9783742 100644 --- a/tests/test_file_size.py +++ b/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 diff --git a/tests/test_query_device_capacity.py b/tests/test_query_device_capacity.py index 374a772..5c0ff35 100644 --- a/tests/test_query_device_capacity.py +++ b/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)