Merge pull request #126 from timlnx/fixup-query-capacity-stuff

Reworking query_*_capacity. Closes #125

Author: timlnx

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`
 

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:

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
    <https://github.com/timlnx/bitmath/issues/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
    <https://github.com/timlnx/bitmath/issues/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

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 <sys/disk.h>: 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 <sys/disk.h>: 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