5252import os
5353import os .path
5454import platform
55+ import re
56+ import shutil
5557import sys
5658import threading
5759
5860from collections .abc import Generator , Iterable , Iterator
59- from typing import IO , Any
61+ from typing import IO , Any , NamedTuple , Union
6062
6163# For device capacity reading in query_device_capacity().
6264if os .name == 'posix' :
6971 import msvcrt
7072
7173#: Platforms where :func:`query_device_capacity` is supported.
72- #: Corresponds to possible values of :data:`os.name`.
74+ #: Corresponds to possible values of :data:`os.name`. macOS (Darwin)
75+ #: is not supported due to SIP restrictions on raw block device access.
7376SUPPORTED_PLATFORMS = frozenset ({'posix' , 'nt' })
7477
7578__all__ = ['Bit' , 'Byte' , 'KiB' , 'MiB' , 'GiB' , 'TiB' , 'PiB' , 'EiB' , 'ZiB' , 'YiB' ,
7881 'Pb' , 'Eb' , 'Zb' , 'Yb' , 'getsize' , 'listdir' , 'format' ,
7982 'format_string' , 'format_plural' , 'parse_string' , 'parse_string_unsafe' ,
8083 'sum' , 'ALL_UNIT_TYPES' , 'NIST' , 'NIST_PREFIXES' , 'NIST_STEPS' ,
81- 'SI' , 'SI_PREFIXES' , 'SI_STEPS' ]
84+ 'SI' , 'SI_PREFIXES' , 'SI_STEPS' , 'Capacity' , 'query_capacity' ,
85+ 'query_device_capacity' ]
8286
8387#: A list of all the valid prefix unit types. Mostly for reference,
8488#: also used by the CLI tool as valid types
@@ -1365,31 +1369,33 @@ class DISK_GEOMETRY_EX(ctypes.Structure):
13651369
13661370
13671371def query_device_capacity (device_fd : IO [Any ]) -> Byte :
1368- """Create bitmath instances of the capacity of a system block device
1372+ """Query the raw physical capacity of a block device.
13691373
1370- Make one or more ioctl request to query the capacity of a block
1371- device. Perform any processing required to compute the final capacity
1372- value. Return the device capacity in bytes as a :class:`bitmath.Byte`
1373- instance.
1374-
1375- Thanks to the following resources for help figuring this out Linux/Mac
1376- ioctl's for querying block device sizes:
1374+ Most users should prefer :func:`query_capacity`. This function is for
1375+ callers who need raw physical device capacity (e.g. disk imaging tools).
1376+ Requires root on Linux and administrator on Windows. Not supported on
1377+ macOS (SIP restriction).
13771378
1378- * http://stackoverflow.com/a/12925285/263969
1379- * http://stackoverflow.com/a/9764508/263969
1380-
1381- :param file device_fd: A ``file`` object of the device to query the
1382- capacity of. On Linux/macOS: ``open("/dev/sda", "rb")``. On Windows:
1383- ``open(r'\\ \\ .\\ PhysicalDrive0', 'rb')`` (requires administrator privileges).
1379+ :param file device_fd: A ``file`` object of the device to query.
1380+ On Linux: ``open("/dev/sda", "rb")`` (requires root).
1381+ On Windows: ``open(r'\\ \\ .\\ PhysicalDrive0', 'rb')`` (requires administrator).
13841382
13851383 :return: a bitmath :class:`bitmath.Byte` instance equivalent to the
13861384 capacity of the target device in bytes.
1385+ :raises NotImplementedError: on macOS or any other unsupported platform.
1386+ :raises ValueError: if the file descriptor is not a block device.
13871387"""
13881388 if os .name not in SUPPORTED_PLATFORMS :
13891389 raise NotImplementedError (f"'bitmath.query_device_capacity' is not supported on this platform: { os .name } " )
13901390 if os .name == 'nt' :
13911391 return Byte (_query_device_capacity_windows (device_fd ))
13921392
1393+ if platform .system () == 'Darwin' :
1394+ raise NotImplementedError (
1395+ "query_device_capacity is not supported on macOS; "
1396+ "SIP blocks raw block device access. Use query_capacity() instead."
1397+ )
1398+
13931399 s = os .stat (device_fd .name ).st_mode
13941400 if not stat .S_ISBLK (s ):
13951401 raise ValueError ("The file descriptor provided is not of a device type" )
@@ -1434,45 +1440,6 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte:
14341440 # BLKGETSIZE64.
14351441 "func" : lambda x : x ["BLKGETSIZE64" ]
14361442 },
1437- # ioctls for the "Darwin" (Mac OS X) platform
1438- "Darwin" : {
1439- "request_params" : [
1440- # A list of parameters to calculate the block size.
1441- #
1442- # ( PARAM_NAME , FORMAT_CHAR , REQUEST_CODE )
1443- ("DKIOCGETBLOCKCOUNT" , "L" , 0x40086419 ),
1444- # Per <sys/disk.h>: get media's block count - uint64_t
1445- #
1446- # As in the BLKGETSIZE64 example, an unsigned 64 bit
1447- # integer will use the 'L' formatting character
1448- ("DKIOCGETBLOCKSIZE" , "I" , 0x40046418 )
1449- # Per <sys/disk.h>: get media's block size - uint32_t
1450- #
1451- # This request returns an unsigned 32 bit integer, or
1452- # in other words: just a normal integer (or 'int' c
1453- # type). That should require 4 bytes of space for
1454- # buffering. According to the struct modules
1455- # 'Formatting Characters' chart:
1456- #
1457- # * Character 'I' - Unsigned Int C Type (uint32_t) - Loads into a Python int type
1458- ],
1459- # OS X doesn't have a direct equivalent to the Linux
1460- # BLKGETSIZE64 request. Instead, we must request how many
1461- # blocks (or "sectors") are on the disk, and the size (in
1462- # bytes) of each block. Finally, multiply the two together
1463- # to obtain capacity:
1464- #
1465- # n Block * y Byte
1466- # capacity (bytes) = -------
1467- # 1 Block
1468- "func" : lambda x : x ["DKIOCGETBLOCKCOUNT" ] * x ["DKIOCGETBLOCKSIZE" ]
1469- # This expression simply accepts a dictionary ``x`` as a
1470- # parameter, and then returns the result of multiplying
1471- # the two named dictionary items together. In this case,
1472- # that means multiplying ``DKIOCGETBLOCKCOUNT``, the total
1473- # number of blocks, by ``DKIOCGETBLOCKSIZE``, the size of
1474- # each block in bytes.
1475- }
14761443 }
14771444
14781445 platform_params = ioctl_map [platform .system ()]
@@ -1486,7 +1453,7 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte:
14861453 # conditions for some possible errors. Really only for cases
14871454 # where it would add value to override the default exception
14881455 # message string.
1489- buffer = fcntl .ioctl (device_fd .fileno (), request_code , buffer_size )
1456+ buffer = fcntl .ioctl (device_fd .fileno (), request_code , b' \x00 ' * buffer_size )
14901457
14911458 # Unpack the raw result from the ioctl call into a familiar
14921459 # python data type according to the ``fmt`` rules.
@@ -1497,6 +1464,71 @@ def query_device_capacity(device_fd: IO[Any]) -> Byte:
14971464 return Byte (platform_params ['func' ](results ))
14981465
14991466
1467+ class Capacity (NamedTuple ):
1468+ """Capacity of a filesystem volume returned by :func:`query_capacity`."""
1469+ total : 'Bitmath'
1470+ used : 'Bitmath'
1471+ free : 'Bitmath'
1472+
1473+
1474+ # Matches a bare drive letter: "C", "c", "C:", "c:" — nothing else.
1475+ _DRIVE_LETTER_RE = re .compile (r'^[A-Za-z]:?$' )
1476+
1477+
1478+ def query_capacity (path : Union [str , os .PathLike ], bestprefix : bool = True ,
1479+ system : int = NIST ) -> Capacity :
1480+ """Return the total, used, and free capacity of the volume at ``path``.
1481+
1482+ This is the recommended API for querying volume or mount-point size. It
1483+ works cross-platform without elevated privileges.
1484+
1485+ :param path: A path on the filesystem volume to query. On Windows, a
1486+ bare drive letter (``"C"``, ``"C:"``) is normalized to ``"C:\\ "``.
1487+ :param bool bestprefix: When ``True`` (default), each field of the
1488+ returned :class:`Capacity` is already normalized via
1489+ :meth:`~bitmath.Bitmath.best_prefix` for human-readable output.
1490+ When ``False``, each field is a raw :class:`bitmath.Byte`.
1491+ :param int system: Unit system to use when ``bestprefix`` is ``True``.
1492+ Either :data:`bitmath.NIST` (default, binary prefixes like ``GiB``)
1493+ or :data:`bitmath.SI` (decimal prefixes like ``GB``). Ignored when
1494+ ``bestprefix`` is ``False``.
1495+
1496+ :return: A :class:`Capacity` NamedTuple with ``total``, ``used``, and
1497+ ``free`` fields, each a :class:`bitmath.Bitmath` instance.
1498+
1499+ :raises FileNotFoundError: if ``path`` does not exist.
1500+ :raises PermissionError: if the process lacks access to query ``path``.
1501+
1502+ Example — attribute access (human-readable by default)::
1503+
1504+ cap = bitmath.query_capacity("/")
1505+ print(cap.total) # e.g. 465.762 GiB
1506+
1507+ Example — tuple unpacking::
1508+
1509+ total, used, free = bitmath.query_capacity("/")
1510+
1511+ Example — raw bytes and SI prefixes::
1512+
1513+ cap_raw = bitmath.query_capacity("/", bestprefix=False)
1514+ cap_si = bitmath.query_capacity("/", system=bitmath.SI)
1515+ """
1516+ normalized : Union [str , os .PathLike ] = path
1517+ if os .name == 'nt' :
1518+ s = str (path ).upper ()
1519+ if _DRIVE_LETTER_RE .match (s ):
1520+ normalized = s .rstrip (':' ) + ':\\ '
1521+ usage = shutil .disk_usage (normalized )
1522+ total , used , free = Byte (usage .total ), Byte (usage .used ), Byte (usage .free )
1523+ if bestprefix :
1524+ return Capacity (
1525+ total .best_prefix (system = system ),
1526+ used .best_prefix (system = system ),
1527+ free .best_prefix (system = system ),
1528+ )
1529+ return Capacity (total , used , free )
1530+
1531+
15001532def getsize (path : str , bestprefix : bool = True , system : int = NIST ) -> Bitmath :
15011533 """Return a bitmath instance in the best human-readable representation
15021534of the file size at `path`. Optionally, provide a preferred unit
0 commit comments