"""Utility functions for formatting and parsing.""" import pathlib import re import sys __all__ = [ "get_output_size", "get_sector_size", "parse_size", "sparse_range", ] # Cache for sector size lookup (path -> size) _sector_size_cache: dict[str, int] = {} def get_sector_size(path: str | pathlib.Path) -> int: """Get sector size for a block device, or 512 as fallback.""" import os import stat try: st = pathlib.Path(path).stat() if not stat.S_ISBLK(st.st_mode): return 512 except OSError: return 512 try: fd = os.open(str(path), os.O_RDONLY) try: import fcntl import struct if sys.platform == "darwin": # macOS: DKIOCGETBLOCKSIZE = 0x40046418 DKIOCGETBLOCKSIZE = 0x40046418 buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4) return struct.unpack("I", buf)[0] else: # Linux: BLKSSZGET = 0x1268 BLKSSZGET = 0x1268 buf = fcntl.ioctl(fd, BLKSSZGET, b"\x00" * 4) return struct.unpack("i", buf)[0] finally: os.close(fd) except (OSError, ImportError, Exception): return 512 def get_output_size(path: str | pathlib.Path | None) -> int | None: """Get the size of output file or block device. Returns None for stdout or non-existent files. Returns the size in bytes for existing files or block devices. """ import os import stat if not path: return None # stdout try: st = pathlib.Path(path).stat() except OSError: return None # doesn't exist yet if stat.S_ISBLK(st.st_mode): # Block device - get size via ioctl try: fd = os.open(str(path), os.O_RDONLY) try: import fcntl import struct if sys.platform == "darwin": # macOS: DKIOCGETBLOCKCOUNT and DKIOCGETBLOCKSIZE DKIOCGETBLOCKCOUNT = 0x40086419 DKIOCGETBLOCKSIZE = 0x40046418 count_buf = fcntl.ioctl(fd, DKIOCGETBLOCKCOUNT, b"\x00" * 8) size_buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4) block_count = struct.unpack("Q", count_buf)[0] block_size = struct.unpack("I", size_buf)[0] return block_count * block_size else: # Linux: BLKGETSIZE64 = 0x80081272 BLKGETSIZE64 = 0x80081272 buf = fcntl.ioctl(fd, BLKGETSIZE64, b"\x00" * 8) return struct.unpack("Q", buf)[0] finally: os.close(fd) except (OSError, ImportError, Exception): return None elif stat.S_ISREG(st.st_mode): return st.st_size else: return None def parse_size(length: str | None, output_path: str | pathlib.Path | None = None) -> int | None: """Parse size string with SI/IEC prefixes. Supports: - Plain numbers: 1000, 1_000_000 - SI prefixes: k, m, g, t, p (powers of 1000) - IEC prefixes: ki, mi, gi, ti, pi (powers of 1024) - Optional 'b' suffix: kb, kib, mb, mib, etc. - Special unit 'sect' = device sector size (detected, fallback 512) - Case insensitive Examples: 1k, 1ki, 1kb, 1kib, 100m, 100mi, 1g, 1gi, 10sect """ if length is None: return None s = length.strip().lower().replace("_", "") # Handle sect unit m = re.match(r"^(\d+)\s*sect?s?$", s) if m: if output_path: path_str = str(output_path) if path_str not in _sector_size_cache: _sector_size_cache[path_str] = get_sector_size(output_path) sector_size = _sector_size_cache.get(path_str, 512) else: sector_size = 512 return int(m.group(1)) * sector_size # SI/IEC prefixes si_prefixes = {"k": 1000, "m": 1000**2, "g": 1000**3, "t": 1000**4, "p": 1000**5} iec_prefixes = { "ki": 1024, "mi": 1024**2, "gi": 1024**3, "ti": 1024**4, "pi": 1024**5, } # Try IEC first (ki, mi, etc.) - must check before SI m = re.match(r"^(\d+(?:\.\d+)?)\s*(ki|mi|gi|ti|pi)b?$", s) if m: num, prefix = m.groups() return int(float(num) * iec_prefixes[prefix]) # Try SI (k, m, g, etc.) m = re.match(r"^(\d+(?:\.\d+)?)\s*([kmgtp])b?$", s) if m: num, prefix = m.groups() return int(float(num) * si_prefixes[prefix]) # Plain number m = re.match(r"^(\d+)$", s) if m: return int(m.group(1)) raise ValueError(f"Invalid size format: {length}") def sparse_range(n: int, max_items: int = 9) -> list[int]: """Generate a sparse range from 1 to N for benchmarking thread counts.""" if n < 1: return [1] if n <= max_items - 1: return list(range(n + 1)) keep = 3 # dense prefix: 1,2,3 out = list(range(keep + 1)) remaining = max_items - keep step = max(1, n // (remaining - 1)) for k in range(1, remaining): v = k * step if v > out[-1]: out.append(v) if out[-1] != n: out[-1] = n return out