qq_lib.properties.size

  1# Released under MIT License.
  2# Copyright (c) 2025-2026 Ladislav Bartos and Robert Vacha Lab
  3
  4"""
  5Utility class for representing and manipulating memory and storage sizes.
  6
  7This module defines the `Size` class, a numeric wrapper used across
  8qq to express quantities such as memory limits and scratch allocations.
  9"""
 10
 11import math
 12import re
 13from dataclasses import dataclass
 14from typing import Self
 15
 16from qq_lib.core.config import CFG
 17from qq_lib.core.error import QQError
 18
 19
 20@dataclass(init=False)
 21class Size:
 22    """
 23    Represents a memory or storage size.
 24
 25    The value is stored internally in kilobytes (kB). When converted to a string,
 26    it is displayed in the largest human-readable unit such that the relative
 27    rounding error does not exceed `CFG.size.max_rounding_error`.
 28    """
 29
 30    value: int
 31
 32    _unit_map = {
 33        "kb": 1,
 34        "mb": 1024,
 35        "gb": 1024 * 1024,
 36        "tb": 1024 * 1024 * 1024,
 37        "pb": 1024 * 1024 * 1024 * 1024,
 38    }
 39
 40    def __init__(self, value: int, unit: str = "kb"):
 41        unit = unit.lower()
 42        if unit not in self._unit_map:
 43            # special handling of bytes
 44            if unit == "b":
 45                self.value = 0 if value == 0 else 1
 46                return
 47
 48            raise QQError(f"Unsupported unit for size '{unit}'.")
 49
 50        self.value = value * self._unit_map[unit]
 51
 52    @classmethod
 53    def from_string(cls, s: str) -> Self:
 54        """
 55        Create a Size object from a string.
 56
 57        Args:
 58            s (str): A string representation of the size, e.g., "10mb", "10 mb", "10m", "10M".
 59
 60        Returns:
 61            Size: A Size instance with parsed value and unit.
 62
 63        Raises:
 64            QQError: If the string cannot be parsed or contains an invalid unit.
 65        """
 66        match = re.match(r"^\s*(\d+)\s*([a-zA-Z]+)\s*$", s)
 67        if not match:
 68            raise QQError(f"Invalid size string: '{s}'.")
 69        value, unit = match.groups()
 70
 71        # normalize single-letter units to their full form by appending 'b'
 72        # but skip bytes
 73        if len(unit) == 1 and unit != "b":
 74            unit = unit.lower() + "b"
 75
 76        return cls(int(value), unit)
 77
 78    def __mul__(self, n: int) -> "Size":
 79        """
 80        Multiply the Size by an integer.
 81
 82        Args:
 83            n (int): The multiplier.
 84
 85        Returns:
 86            Size: A new Size object with the scaled value.
 87
 88        Raises:
 89            TypeError: If the multiplier is not an integer.
 90        """
 91        if not isinstance(n, int):
 92            return NotImplemented
 93        if n == 0:
 94            return Size(0, "kb")
 95
 96        return Size(self.value * n, "kb")
 97
 98    # allow 3 * Size
 99    __rmul__ = __mul__
100
101    def __str__(self) -> str:
102        for unit, factor in reversed(list(self._unit_map.items())):
103            value = self.value / factor
104
105            if value >= 1:
106                rounded = round(value)
107                # compute relative error from rounding
108                approx_kb = rounded * factor
109                error = abs(approx_kb - self.value) / self.value
110                if error <= CFG.size.max_rounding_error or unit == "kb":
111                    return f"{rounded}{unit}"
112                # otherwise, try smaller unit
113
114        # should not get here
115        return f"{self.value}kb"
116
117    def to_str_exact(self) -> str:
118        """Convert the Size to string while keeping it in kilobytes."""
119        return f"{self.value}kb"
120
121    def to_str_exact_slurm(self) -> str:
122        """Convert the Size to string while keeping it in kilobytes. Use K for the unit."""
123        return f"{self.value}K"
124
125    def __floordiv__(self, n: int) -> "Size":
126        """
127        Divide the Size by an integer.
128
129        Args:
130            n (int): The divisor.
131
132        Returns:
133            Size: A new Size object representing the divided size.
134
135        Raises:
136            TypeError: If n is not an integer.
137            ZeroDivisionError: If n is zero.
138        """
139        if not isinstance(n, int):
140            return NotImplemented
141        if n == 0:
142            raise ZeroDivisionError("Division by zero.")
143
144        return Size(math.ceil(self.value / n), "kb")
145
146    def __truediv__(self, other: "Size") -> float:
147        """
148        Perform true division (/) between two Size instances.
149
150        Computes the ratio of this Size to another, expressed as a float.
151
152        Args:
153            other (Size): The divisor Size instance.
154
155        Returns:
156            float: The ratio of self to other, based on total kilobytes.
157
158        Raises:
159            TypeError:
160                If `other` is not a Size instance.
161            ZeroDivisionError:
162                If `other` is a zero Size.
163        """
164        if not isinstance(other, Size):
165            raise TypeError(
166                f"Unsupported operand type(s) for /: 'Size' and '{type(other).__name__}'"
167            )
168
169        if other.value == 0:
170            raise ZeroDivisionError("Division by zero size.")
171
172        return self.value / other.value
173
174    def __sub__(self, other: "Size") -> "Size":
175        """
176        Subtract one Size from another.
177
178        Args:
179            other (Size): The Size instance to subtract.
180
181        Returns:
182            Size: A new Size instance representing the difference.
183
184        Raises:
185            TypeError: If `other` is not a Size instance.
186            ValueError: If the result would be negative.
187        """
188        if not isinstance(other, Size):
189            raise TypeError(
190                f"Unsupported operand type(s) for -: 'Size' and '{type(other).__name__}'"
191            )
192
193        result_kb = self.value - other.value
194        if result_kb < 0:
195            raise ValueError("Resulting Size cannot be negative.")
196
197        return Size(result_kb, "kb")