diff --git a/CLAUDE.md b/CLAUDE.md index 9b39122..ca9b42f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -61,7 +61,8 @@ All unit values are normalized to bits internally; conversion between units happ - `getsize(path, ...)` — file size with automatic prefix selection - `listdir(search_base, ...)` — recursive directory listing with sizes - `parse_string(s)` / `parse_string_unsafe(s, system=SI)` — string → bitmath object -- `query_device_capacity(device_fd)` — POSIX device capacity (Linux/macOS) +- `query_capacity(path)` — volume/mount-point capacity as a `Capacity(total, used, free)` NamedTuple of `Byte` instances; cross-platform, no elevated privileges required +- `query_device_capacity(device_fd)` — raw physical block device capacity (Linux: requires root; Windows: requires administrator; macOS: raises `NotImplementedError` due to SIP) **Constants:** `NIST`, `SI`, `NIST_PREFIXES`, `SI_PREFIXES`, `ALL_UNIT_TYPES` diff --git a/Makefile b/Makefile index f94b268..e505b78 100644 --- a/Makefile +++ b/Makefile @@ -112,8 +112,10 @@ pypitest: build @echo "#############################################" . $(NAME)env3/bin/activate && pip install twine && twine upload --repository testpypi dist/* -# usage example: make tag TAG=1.1.0-1 +# usage example: make tag TAG=v1.1.0-1 tag: + @if [ -z "$(TAG)" ]; then echo "ERROR: TAG is required. Example: make tag TAG=v2.0.0"; exit 1; fi + @case "$(TAG)" in v*) ;; *) echo "ERROR: TAG must start with 'v'. Got: '$(TAG)'. Example: make tag TAG=v2.0.0"; exit 1 ;; esac git tag -s -m $(TAG) $(TAG) clean: diff --git a/NEWS.rst b/NEWS.rst index f1132b3..ea68be2 100644 --- a/NEWS.rst +++ b/NEWS.rst @@ -54,6 +54,13 @@ Breaking Changes ``"Byte"`` or ``"Bit"`` will need to be updated. The class names themselves are unchanged. +**query_device_capacity() on macOS** + :func:`bitmath.query_device_capacity` now raises + :exc:`NotImplementedError` on macOS. System Integrity Protection + (SIP) blocks raw block device access even for root, making the + previous ioctl path unreliable. Use :func:`bitmath.query_capacity` + instead. + **Build and install** ``setup.py`` and ``setup.py.in`` are gone. Installation is ``pip install bitmath``. Source builds use ``python -m build``. @@ -101,6 +108,17 @@ still works exactly the same way. What 2.0.0 adds on top of that: family is now preserved. Closes `issue #95 `_. +**query_capacity() — recommended volume size API** + New :func:`bitmath.query_capacity` returns a :class:`bitmath.Capacity` + NamedTuple of ``(total, used, free)`` :class:`bitmath.Bitmath` + instances for any path or mount point. Works cross-platform without + elevated privileges. This is the recommended API for "how big is + this volume?" queries. Accepts ``bestprefix=True`` (default) to get + already human-readable units (e.g. ``GiB``) and ``system=bitmath.SI`` + to opt into decimal prefixes instead of the default NIST binary + prefixes. Set ``bestprefix=False`` to receive raw + :class:`bitmath.Byte` instances. + **Windows device capacity** :func:`bitmath.query_device_capacity` now works on Windows via ``DeviceIoControl``. Open the device as @@ -110,6 +128,13 @@ still works exactly the same way. What 2.0.0 adds on top of that: platforms where the function is available. Closes `issue #52 `_. +**query_device_capacity() Linux buffer fix** + :func:`bitmath.query_device_capacity` on Linux was passing an + integer where an ioctl buffer was required, causing + ``OSError: [Errno 14] Bad address``. The call now correctly + allocates a zero-filled buffer of the proper size via + ``b'\\x00' * struct.calcsize(fmt)``. + **Flexible string parsing** :func:`bitmath.parse_string` with ``strict=False`` accepts ambiguous input such as ``"1g"`` or ``"1GB"`` and resolves it to diff --git a/bitmath/__init__.py b/bitmath/__init__.py index 3739e1c..6fe98da 100644 --- a/bitmath/__init__.py +++ b/bitmath/__init__.py @@ -52,11 +52,13 @@ import os import os.path import platform +import re +import shutil import sys import threading from collections.abc import Generator, Iterable, Iterator -from typing import IO, Any +from typing import IO, Any, NamedTuple, Union # For device capacity reading in query_device_capacity(). if os.name == 'posix': @@ -69,7 +71,8 @@ import msvcrt #: Platforms where :func:`query_device_capacity` is supported. -#: Corresponds to possible values of :data:`os.name`. +#: Corresponds to possible values of :data:`os.name`. macOS (Darwin) +#: is not supported due to SIP restrictions on raw block device access. SUPPORTED_PLATFORMS = frozenset({'posix', 'nt'}) __all__ = ['Bit', 'Byte', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB', @@ -78,7 +81,8 @@ 'Pb', 'Eb', 'Zb', 'Yb', 'getsize', 'listdir', 'format', 'format_string', 'format_plural', 'parse_string', 'parse_string_unsafe', 'sum', 'ALL_UNIT_TYPES', 'NIST', 'NIST_PREFIXES', 'NIST_STEPS', - 'SI', 'SI_PREFIXES', 'SI_STEPS'] + 'SI', 'SI_PREFIXES', 'SI_STEPS', 'Capacity', 'query_capacity', + 'query_device_capacity'] #: A list of all the valid prefix unit types. Mostly for reference, #: also used by the CLI tool as valid types @@ -1342,31 +1346,33 @@ class DISK_GEOMETRY_EX(ctypes.Structure): def query_device_capacity(device_fd: IO[Any]) -> Byte: - """Create bitmath instances of the capacity of a system block device + """Query the raw physical capacity of a block device. -Make one or more ioctl request to query the capacity of a block -device. Perform any processing required to compute the final capacity -value. Return the device capacity in bytes as a :class:`bitmath.Byte` -instance. - -Thanks to the following resources for help figuring this out Linux/Mac -ioctl's for querying block device sizes: +Most users should prefer :func:`query_capacity`. This function is for +callers who need raw physical device capacity (e.g. disk imaging tools). +Requires root on Linux and administrator on Windows. Not supported on +macOS (SIP restriction). -* http://stackoverflow.com/a/12925285/263969 -* http://stackoverflow.com/a/9764508/263969 - - :param file device_fd: A ``file`` object of the device to query the - capacity of. On Linux/macOS: ``open("/dev/sda", "rb")``. On Windows: - ``open(r'\\\\.\\PhysicalDrive0', 'rb')`` (requires administrator privileges). + :param file device_fd: A ``file`` object of the device to query. + On Linux: ``open("/dev/sda", "rb")`` (requires root). + On Windows: ``open(r'\\\\.\\PhysicalDrive0', 'rb')`` (requires administrator). :return: a bitmath :class:`bitmath.Byte` instance equivalent to the capacity of the target device in bytes. + :raises NotImplementedError: on macOS or any other unsupported platform. + :raises ValueError: if the file descriptor is not a block device. """ 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)) + if platform.system() == 'Darwin': + raise NotImplementedError( + "query_device_capacity is not supported on macOS; " + "SIP blocks raw block device access. Use query_capacity() instead." + ) + s = os.stat(device_fd.name).st_mode if not stat.S_ISBLK(s): raise ValueError("The file descriptor provided is not of a device type") @@ -1411,45 +1417,6 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte: # BLKGETSIZE64. "func": lambda x: x["BLKGETSIZE64"] }, - # ioctls for the "Darwin" (Mac OS X) platform - "Darwin": { - "request_params": [ - # A list of parameters to calculate the block size. - # - # ( PARAM_NAME , FORMAT_CHAR , REQUEST_CODE ) - ("DKIOCGETBLOCKCOUNT", "L", 0x40086419), - # Per : get media's block count - uint64_t - # - # As in the BLKGETSIZE64 example, an unsigned 64 bit - # integer will use the 'L' formatting character - ("DKIOCGETBLOCKSIZE", "I", 0x40046418) - # Per : get media's block size - uint32_t - # - # This request returns an unsigned 32 bit integer, or - # in other words: just a normal integer (or 'int' c - # type). That should require 4 bytes of space for - # buffering. According to the struct modules - # 'Formatting Characters' chart: - # - # * Character 'I' - Unsigned Int C Type (uint32_t) - Loads into a Python int type - ], - # OS X doesn't have a direct equivalent to the Linux - # BLKGETSIZE64 request. Instead, we must request how many - # blocks (or "sectors") are on the disk, and the size (in - # bytes) of each block. Finally, multiply the two together - # to obtain capacity: - # - # n Block * y Byte - # capacity (bytes) = ------- - # 1 Block - "func": lambda x: x["DKIOCGETBLOCKCOUNT"] * x["DKIOCGETBLOCKSIZE"] - # This expression simply accepts a dictionary ``x`` as a - # parameter, and then returns the result of multiplying - # the two named dictionary items together. In this case, - # that means multiplying ``DKIOCGETBLOCKCOUNT``, the total - # number of blocks, by ``DKIOCGETBLOCKSIZE``, the size of - # each block in bytes. - } } platform_params = ioctl_map[platform.system()] @@ -1463,7 +1430,7 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte: # conditions for some possible errors. Really only for cases # where it would add value to override the default exception # message string. - buffer = fcntl.ioctl(device_fd.fileno(), request_code, buffer_size) + buffer = fcntl.ioctl(device_fd.fileno(), request_code, b'\x00' * buffer_size) # Unpack the raw result from the ioctl call into a familiar # python data type according to the ``fmt`` rules. @@ -1474,6 +1441,71 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte: return Byte(platform_params['func'](results)) +class Capacity(NamedTuple): + """Capacity of a filesystem volume returned by :func:`query_capacity`.""" + total: 'Bitmath' + used: 'Bitmath' + free: 'Bitmath' + + +# Matches a bare drive letter: "C", "c", "C:", "c:" — nothing else. +_DRIVE_LETTER_RE = re.compile(r'^[A-Za-z]:?$') + + +def query_capacity(path: Union[str, os.PathLike], bestprefix: bool = True, + system: int = NIST) -> Capacity: + """Return the total, used, and free capacity of the volume at ``path``. + +This is the recommended API for querying volume or mount-point size. It +works cross-platform without elevated privileges. + + :param path: A path on the filesystem volume to query. On Windows, a + bare drive letter (``"C"``, ``"C:"``) is normalized to ``"C:\\"``. + :param bool bestprefix: When ``True`` (default), each field of the + returned :class:`Capacity` is already normalized via + :meth:`~bitmath.Bitmath.best_prefix` for human-readable output. + When ``False``, each field is a raw :class:`bitmath.Byte`. + :param int system: Unit system to use when ``bestprefix`` is ``True``. + Either :data:`bitmath.NIST` (default, binary prefixes like ``GiB``) + or :data:`bitmath.SI` (decimal prefixes like ``GB``). Ignored when + ``bestprefix`` is ``False``. + + :return: A :class:`Capacity` NamedTuple with ``total``, ``used``, and + ``free`` fields, each a :class:`bitmath.Bitmath` instance. + + :raises FileNotFoundError: if ``path`` does not exist. + :raises PermissionError: if the process lacks access to query ``path``. + +Example — attribute access (human-readable by default):: + + cap = bitmath.query_capacity("/") + print(cap.total) # e.g. 465.762 GiB + +Example — tuple unpacking:: + + total, used, free = bitmath.query_capacity("/") + +Example — raw bytes and SI prefixes:: + + cap_raw = bitmath.query_capacity("/", bestprefix=False) + cap_si = bitmath.query_capacity("/", system=bitmath.SI) +""" + normalized: Union[str, os.PathLike] = path + if os.name == 'nt': + s = str(path).upper() + if _DRIVE_LETTER_RE.match(s): + normalized = s.rstrip(':') + ':\\' + usage = shutil.disk_usage(normalized) + total, used, free = Byte(usage.total), Byte(usage.used), Byte(usage.free) + if bestprefix: + return Capacity( + total.best_prefix(system=system), + used.best_prefix(system=system), + free.best_prefix(system=system), + ) + return Capacity(total, used, free) + + def getsize(path: str, bestprefix: bool = True, system: int = NIST) -> Bitmath: """Return a bitmath instance in the best human-readable representation of the file size at `path`. Optionally, provide a preferred unit diff --git a/docsite/source/module.rst b/docsite/source/module.rst index a0625c7..9a0aa04 100644 --- a/docsite/source/module.rst +++ b/docsite/source/module.rst @@ -560,61 +560,151 @@ bitmath.parse_string_unsafe() .. versionadded:: 1.3.1 -bitmath.query_device_capacity() -=============================== +bitmath.query_capacity() +======================== + +.. function:: query_capacity(path, bestprefix=True, system=NIST) + + Return the total, used, and free capacity of the volume containing + ``path`` as a :class:`Capacity` NamedTuple of :class:`bitmath.Bitmath` + instances. This is the recommended API for volume size queries — it + works cross-platform without elevated privileges. + + :param path: A path on the volume to query (``str`` or + :class:`os.PathLike`). On Windows, a bare drive letter + such as ``"C"`` or ``"C:"`` is automatically normalized + to ``"C:\\"``. + :param bool bestprefix: When ``True`` (default), each returned field + is pre-normalized via + :meth:`~bitmath.Bitmath.best_prefix` for + human-readable output. When ``False``, raw + :class:`bitmath.Byte` instances are returned. + :param int system: Prefix system used when ``bestprefix`` is + ``True``. :data:`bitmath.NIST` (default) gives + binary prefixes (``GiB``); :data:`bitmath.SI` + gives decimal prefixes (``GB``). Ignored when + ``bestprefix`` is ``False``. + :return: A :class:`bitmath.Capacity` with fields ``total``, + ``used``, and ``free``. + :raises FileNotFoundError: if ``path`` does not exist. + :raises PermissionError: if the process lacks access to query ``path``. + + Linux examples (human-readable by default): -.. function:: query_device_capacity(device_fd) + .. code-block:: python + + >>> import bitmath + >>> cap = bitmath.query_capacity("/") + >>> print(cap.total) + 465.762 GiB + >>> cap = bitmath.query_capacity("/home") + >>> print(cap.free) + 120.5 GiB + + macOS examples: + + .. code-block:: python + + >>> cap = bitmath.query_capacity("/") + >>> print(cap.total) + 500.068 GiB + >>> cap = bitmath.query_capacity("/Volumes/SomeVolume") + >>> print(cap.used) + 14.311 GiB + + Windows examples (bare drive letters are normalized automatically): + + .. code-block:: python + + >>> cap = bitmath.query_capacity("C:\\") + >>> print(cap.total) + 476.837 GiB + >>> cap = bitmath.query_capacity("C") + >>> print(cap.free) + 200.0 GiB + + Tuple unpacking: + + .. code-block:: python - Create :class:`bitmath.Byte` instances representing the capacity of - a block device. + >>> total, used, free = bitmath.query_capacity("/") + >>> print(total, used, free) + 465.762 GiB 186.264 GiB 279.397 GiB - :param file device_fd: An open file handle (``handle = - open('/dev/sda')``) of the target device. - :return: A :class:`bitmath.Byte` equal to the size of ``device_fd``. - :raises ValueError: if file descriptor ``device_fd`` is not of a - device type. - :raises IOError: + Opt into SI (decimal) prefixes: - * :py:exc:`IOError[13]` - If the effective **uid** of this - process does not have access to issue raw commands to block - devices. I.e., this process does not have super-user rights. - * :py:exc:`IOError[2]` - If the device ``device_fd`` points to - does not exist. + .. code-block:: python + + >>> cap = bitmath.query_capacity("/", system=bitmath.SI) + >>> print(cap.total) + 500.068 GB + + Raw :class:`bitmath.Byte` instances (no prefix normalization): + + .. code-block:: python + + >>> cap = bitmath.query_capacity("/", bestprefix=False) + >>> print(cap.total) + 500068036608.0 Byte + + .. versionadded:: 2.0.0 - .. include:: query_device_capacity_warning.rst +bitmath.query_device_capacity() +================================ + +.. function:: query_device_capacity(device_fd) + Query the raw physical capacity of a block device. This is an + advanced, privileged API intended for callers that need the true + hardware capacity (e.g. disk imaging tools). Most users should use + :func:`query_capacity` instead. - .. include:: example_block_devices.rst + Requires root on Linux and administrator on Windows. **Not supported + on macOS** — SIP blocks raw block device access even for root; use + :func:`query_capacity` there. + :param file device_fd: An open file handle of the block device. + :return: A :class:`bitmath.Byte` equal to the device capacity. + :raises ValueError: if ``device_fd`` is not a block device. + :raises NotImplementedError: on macOS or any unsupported platform. + :raises OSError: if the underlying ioctl or DeviceIoControl call fails. - Here's an example using the ``with`` context manager to open a - device and print its capacity with the best-human readable prefix - (line **3**): + Linux example (requires root): .. code-block:: python :linenos: - :emphasize-lines: 3 >>> import bitmath - >>> with open("/dev/sda") as device: + >>> with open("/dev/sda", "rb") as device: ... size = bitmath.query_device_capacity(device).best_prefix() - ... print("Device %s capacity: %s (%s Bytes)" % (device.name, size, size_bytes)) - Device /dev/sda capacity: 238.474937439 GiB (2.56060514304e+11 Bytes) + ... print(f"Device {device.name} capacity: {size}") + Device /dev/sda capacity: 238.475 GiB + Windows example (requires administrator privileges): - :raises NotImplementedError: if called on an unsupported platform. - Supported platforms are Linux, macOS, - and Windows. + .. code-block:: python + :linenos: - .. note:: **Windows usage**: open the device as - ``open(r'\\.\PhysicalDrive0', 'rb')`` (administrator - privileges required). The device path must start with - ``\\.\`` — passing a regular file path raises - :py:exc:`ValueError`. + >>> import bitmath + >>> drives = [r'\\.\PhysicalDrive0', r'\\.\PhysicalDrive1'] + >>> for path in drives: + ... with open(path, 'rb') as drive: + ... size = bitmath.query_device_capacity(drive).best_prefix() + ... print(f"Drive {path}: {size}") + Drive \\.\PhysicalDrive0: 80.0 GiB + Drive \\.\PhysicalDrive1: 14.311 TiB + + .. note:: **macOS**: this function raises :exc:`NotImplementedError` + because System Integrity Protection (SIP) prevents raw + block device access. Use :func:`query_capacity` instead. .. versionadded:: 1.2.4 + .. versionchanged:: 2.0.0 + Added Windows support. macOS now raises :exc:`NotImplementedError` + (SIP restriction). + .. _module_context_managers: Context Managers diff --git a/tests/test_query_capacity.py b/tests/test_query_capacity.py new file mode 100644 index 0000000..a6f8b84 --- /dev/null +++ b/tests/test_query_capacity.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- +# SPDX-License-Identifier: MIT +# The MIT License (MIT) +# +# Copyright © 2026 Tim Case +# +# Permission is hereby granted, free of charge, to any person +# obtaining a copy of this software and associated documentation files +# (the "Software"), to deal in the Software without restriction, +# including without limitation the rights to use, copy, modify, merge, +# publish, distribute, sublicense, and/or sell copies of the Software, +# and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be +# included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS +# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +""" +Tests for bitmath.query_capacity() +""" + +import os +import pathlib +from unittest import skipUnless + +import bitmath +from bitmath import Bitmath, Byte, Capacity + +from . import TestCase + + +class TestQueryCapacity(TestCase): + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_posix_root(self): + """query_capacity('/') returns a valid Capacity on POSIX""" + result = bitmath.query_capacity("/") + self.assertIsInstance(result, Capacity) + self.assertIsInstance(result.total, Bitmath) + self.assertIsInstance(result.used, Bitmath) + self.assertIsInstance(result.free, Bitmath) + self.assertGreater(result.total.bytes, 0) + self.assertLessEqual(result.used.bytes + result.free.bytes, result.total.bytes) + + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_posix_cwd(self): + """query_capacity('.') returns valid Capacity for cwd on POSIX""" + result = bitmath.query_capacity(".") + self.assertIsInstance(result, Capacity) + self.assertGreater(result.total.bytes, 0) + self.assertLessEqual(result.used.bytes + result.free.bytes, result.total.bytes) + + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_tuple_unpacking(self): + """query_capacity result supports tuple unpacking""" + t, u, f = bitmath.query_capacity("/") + self.assertIsInstance(t, Bitmath) + self.assertIsInstance(u, Bitmath) + self.assertIsInstance(f, Bitmath) + + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_bestprefix_false_returns_bytes(self): + """bestprefix=False returns raw Byte instances""" + result = bitmath.query_capacity("/", bestprefix=False) + self.assertIs(type(result.total), Byte) + self.assertIs(type(result.used), Byte) + self.assertIs(type(result.free), Byte) + + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_bestprefix_true_preserves_bytes(self): + """bestprefix=True preserves underlying byte counts""" + raw = bitmath.query_capacity("/", bestprefix=False) + pretty = bitmath.query_capacity("/") + self.assertEqual(raw.total.bytes, pretty.total.bytes) + self.assertEqual(raw.used.bytes, pretty.used.bytes) + self.assertEqual(raw.free.bytes, pretty.free.bytes) + + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_system_nist_vs_si(self): + """system kwarg switches between NIST (GiB) and SI (GB) prefixes""" + nist = bitmath.query_capacity("/", system=bitmath.NIST) + si = bitmath.query_capacity("/", system=bitmath.SI) + # Same underlying bytes, different display units. + self.assertEqual(nist.total.bytes, si.total.bytes) + # NIST units end in 'iB' (e.g. GiB, MiB); SI units do not. + # On a sufficiently large volume these will differ in unit class. + if nist.total.bytes >= 1000: + self.assertTrue(nist.total.unit.endswith('iB') or nist.total.unit == 'Byte') + self.assertFalse(si.total.unit.endswith('iB')) + + @skipUnless(os.name == 'posix', 'POSIX only') + def test_query_capacity_pathlike_input(self): + """query_capacity accepts pathlib.Path input""" + result = bitmath.query_capacity(pathlib.Path("/")) + self.assertIsInstance(result, Capacity) + self.assertGreater(result.total.bytes, 0) + + def test_query_capacity_nonexistent_path_raises(self): + """query_capacity raises FileNotFoundError for a path that does not exist""" + with self.assertRaises((FileNotFoundError, OSError)): + bitmath.query_capacity("/nonexistent/path/that/should/not/exist/xyzzy") + + @skipUnless(os.name == 'nt', 'Windows only') + def test_query_capacity_windows_drive_letter_normalization(self): + """query_capacity normalizes bare drive letters on Windows""" + c_bare = bitmath.query_capacity("C") + c_colon = bitmath.query_capacity("C:") + c_backslash = bitmath.query_capacity("C:\\") + c_lower = bitmath.query_capacity("c") + self.assertEqual(c_bare, c_colon) + self.assertEqual(c_bare, c_backslash) + self.assertEqual(c_bare, c_lower) + + @skipUnless(os.name == 'nt', 'Windows only') + def test_query_capacity_windows_pathlike_drive_letter(self): + """query_capacity normalizes a pathlib drive-letter input on Windows""" + result = bitmath.query_capacity(pathlib.PureWindowsPath("C:")) + self.assertIsInstance(result, Capacity) + self.assertGreater(result.total.bytes, 0) diff --git a/tests/test_query_device_capacity.py b/tests/test_query_device_capacity.py index 5c0ff35..240fbf3 100644 --- a/tests/test_query_device_capacity.py +++ b/tests/test_query_device_capacity.py @@ -78,37 +78,15 @@ def test_query_device_capacity_linux_everything_is_wonderful(self): bytes = bitmath.query_device_capacity(device) self.assertEqual(bytes, 244140625) self.assertEqual(ioctl.call_count, 1) - ioctl.assert_called_once_with(4, 0x80081272, struct.calcsize('L')) + ioctl.assert_called_once_with(4, 0x80081272, b'\x00' * 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( - mock.patch('os.stat'), - mock.patch('stat.S_ISBLK'), - mock.patch('platform.system'), - mock.patch('fcntl.ioctl'), - ) as (os_stat, stat_is_block, plat_system, ioctl): - # These are the struct.pack() equivalents of 244140625 - # (type: u64) and 4096 (type: u32). Multiplied together - # they equal the number of bytes in 1 TB. - returns = [ - struct.pack('L', 244140625), # 'QJ\x8d\x0e\x00\x00\x00\x00' - struct.pack('I', 4096) # , '\x00\x10\x00\x00' - ] - - def side_effect(*args, **kwargs): - return returns.pop(0) - - os_stat.return_value = mock.Mock(st_mode=25008) - stat_is_block.return_value = True - plat_system.return_value = 'Darwin' - ioctl.side_effect = side_effect - - bytes = bitmath.query_device_capacity(device) - # The result should be 1 TB - self.assertEqual(bytes, 1000000000000) - self.assertEqual(ioctl.call_count, 2) + def test_query_device_capacity_macos_raises(self): + """query_device_capacity raises NotImplementedError on macOS (SIP restriction)""" + with mock.patch('bitmath.os.name', 'posix'): + with mock.patch('bitmath.platform.system', return_value='Darwin'): + with self.assertRaises(NotImplementedError) as ctx: + bitmath.query_device_capacity(device) + self.assertIn('SIP', str(ctx.exception)) @skipUnless(os.name == 'posix', 'fcntl is POSIX only') def test_query_device_capacity_device_not_block(self): @@ -116,11 +94,13 @@ def test_query_device_capacity_device_not_block(self): with nested( mock.patch('os.stat'), mock.patch('stat.S_ISBLK'), + mock.patch('platform.system'), mock.patch('fcntl.ioctl'), - ) as (os_stat, stat_is_block, ioctl): + ) as (os_stat, stat_is_block, plat_system, ioctl): os_stat.return_value = mock.Mock(st_mode=33204) # Force ISBLK to reject the input 'device' stat_is_block.return_value = False + plat_system.return_value = 'Linux' with self.assertRaises(ValueError): bitmath.query_device_capacity(non_device_file)