from __future__ import annotations
import re
from typing import IO, Annotated, Any, Literal, Protocol, Union
import numpy as np
import numpy.typing as npt
from typing_extensions import TypeAlias
try:
# Note: Keep import requirements to only stdlib
# as this might be imported from anywhere including
# the bare-install in blender's runtime
import tyro
except ImportError:
tyro = None # type: ignore
_SIZE_SYMBOLS = ("B", "K", "M", "G", "T", "P", "E", "Z", "Y")
_SIZE_BOUNDS = [(1024**i, sym) for i, sym in enumerate(_SIZE_SYMBOLS)]
_SIZE_DICT = {sym: val for val, sym in _SIZE_BOUNDS}
_SIZE_RANGES = list(zip(_SIZE_BOUNDS, _SIZE_BOUNDS[1:]))
FILE_FORMATS = Literal[
"BMP",
"IRIS",
"PNG",
"JPEG",
"JPEG2000",
"TARGA",
"CINEON",
"DPX",
"OPEN_EXR",
"HDR",
"TIFF",
"WEBP",
]
FILE: TypeAlias = Union[None, int, IO[Any]]
EXR_CODECS = Literal["NONE", "PXR24", "ZIP", "PIZ", "RLE", "ZIPS", "DWAA", "DWAB"]
COLOR_MODES = Literal["BW", "RGB", "RGBA"]
def _bytes_from_str(size: str | list[str]) -> int:
# Based on https://stackoverflow.com/a/60708339
if isinstance(size, list):
assert len(size) == 1
return _bytes_from_str(size[0])
try:
return int(size)
except ValueError:
size = size.upper()
if not re.match(r" ", size):
symbols = "".join(_SIZE_SYMBOLS)
size = re.sub(rf"([{symbols}]?B)", r" \1", size)
number, unit = [string.strip() for string in size.split()]
return int(float(number) * _SIZE_DICT[unit[0]])
def _bytes_to_str(nbytes: int, ndigits: int = 1) -> str:
# Based on https://boltons.readthedocs.io/en/latest/strutils.html#boltons.strutils.bytes2human
abs_bytes = abs(nbytes)
for (size, symbol), (next_size, _) in _SIZE_RANGES:
if abs_bytes <= next_size:
break
hnbytes = float(nbytes) / size
return f"{hnbytes:.{ndigits}f}{symbol}"
[docs]
class UpdateFn(Protocol):
"""Mirrors `rich.progress.Progress.update` with curried task-id argument"""
self,
total: float | None = None,
completed: float | None = None,
advance: float | None = None,
description: str | None = None,
visible: bool | None = None,
refresh: bool = False,
) -> None: ...
MemSize: TypeAlias = Annotated[
int,
tyro.constructors.PrimitiveConstructorSpec(
nargs=1,
metavar="BYTES",
instance_from_str=_bytes_from_str,
is_instance=lambda instance: isinstance(instance, int),
str_from_instance=_bytes_to_str,
)
if tyro is not None
else "missing annotation for auto conversion to/from string",
]
_Matrix4x4: TypeAlias = list[list[float]]
Matrix4x4: TypeAlias = Union[_Matrix4x4, npt.NDArray[np.floating]]