238 lines
6.6 KiB
Python
238 lines
6.6 KiB
Python
"""Utility functions for formatting and parsing."""
|
|
|
|
import pathlib
|
|
import re
|
|
import sys
|
|
import time
|
|
|
|
__all__ = [
|
|
"format_size",
|
|
"format_time",
|
|
"get_output_size",
|
|
"get_sector_size",
|
|
"parse_size",
|
|
"print_summary",
|
|
]
|
|
|
|
|
|
def format_size(size: float) -> str:
|
|
"""Format bytes as human-readable size."""
|
|
for unit in ["B", "kB", "MB", "GB", "TB"]:
|
|
if abs(size) < 1000:
|
|
return f"{size:.0f}{unit}"
|
|
size /= 1000
|
|
return f"{size:.0f}PB"
|
|
|
|
|
|
def format_time(seconds: float) -> str:
|
|
"""Format seconds as human-readable time."""
|
|
if seconds < 0:
|
|
return "--"
|
|
if seconds < 1:
|
|
return f"{seconds * 1000:.0f}ms"
|
|
if seconds < 90:
|
|
return f"{seconds:.0f}s"
|
|
elif seconds < 5400: # 90 minutes
|
|
m = int(seconds / 60)
|
|
return f"{m}m"
|
|
elif seconds < 172800: # 48 hours
|
|
h = int(seconds / 3600)
|
|
return f"{h}h"
|
|
else:
|
|
d = int(seconds / 86400)
|
|
return f"{d}d"
|
|
|
|
|
|
def print_summary(written: int, elapsed: float, action: str = "wrote", seed: str | None = None):
|
|
"""Print a nice one-liner summary."""
|
|
speed_gbs = (written / 1_000_000_000) / elapsed if elapsed > 0 else 0
|
|
size_str = format_size(written)
|
|
time_str = format_time(elapsed)
|
|
seed_hint = f" [-s {seed}]" if seed else ""
|
|
print(
|
|
f"RandQuik {action} {size_str} in {time_str} ({speed_gbs:.2f} GB/s){seed_hint}",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
|
|
# 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
|
|
|
|
|
|
def stopwatch():
|
|
t = time.perf_counter()
|
|
|
|
def _generator():
|
|
nonlocal t
|
|
while True:
|
|
t, t0 = time.perf_counter(), t
|
|
yield t - t0
|
|
|
|
return _generator()
|