Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`

Expand Down
4 changes: 3 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
25 changes: 25 additions & 0 deletions NEWS.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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``.
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
146 changes: 89 additions & 57 deletions bitmath/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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':
Expand All @@ -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',
Expand All @@ -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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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()]
Expand All @@ -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.
Expand All @@ -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
Expand Down
Loading
Loading