Files
dopt-basics/src/dopt_basics/batching.py

128 lines
3.8 KiB
Python

from __future__ import annotations
from collections.abc import Collection
from typing import Generic, Never, TypeVar
T = TypeVar("T")
class Chunk(Generic[T]):
def __init__(
self,
capacity: int,
overfill_once: bool = True,
cap_property: str | None = None,
) -> None:
self.capacity = capacity
self.overfill = overfill_once
self.capacity_left: int = capacity
self.capacity_held: int = 0
self.full: bool = False
self._contents: list[T] = []
self.cap_property = cap_property
def __repr__(self) -> str:
return self._contents.__repr__()
def __str__(self) -> str:
return self._contents.__str__()
@property
def contents(self) -> list[T]:
return self._contents
def __len__(self) -> int:
return len(self._contents)
def __getitem__(
self,
idx: int,
) -> T:
return self._contents[idx]
def __raise_for_exceeded_capacity(self) -> Never:
raise ValueError("Item can not be added due to capacity limit.")
def _capacity_sufficient(
self,
size_to_add: int,
) -> bool:
if self.full:
return False
new_size = len(self) + size_to_add
if self.overfill and (new_size > self.capacity):
return True
elif new_size > self.capacity:
return False
else:
return True
def _recalc_capacity(
self,
size_to_add: int,
) -> None:
self.capacity_held += size_to_add
self.capacity_left = max(self.capacity_left - size_to_add, 0)
if self.capacity_left == 0:
self.full = True
else:
self.full = False
def append(
self,
object_: T,
) -> None:
size_to_add: int = 1
if self.cap_property is not None:
if not hasattr(object_, self.cap_property):
raise AttributeError("Object does not posses the wanted property")
attr = getattr(object_, self.cap_property)
if not isinstance(attr, int):
raise TypeError("Capacity property must be an integer")
size_to_add = attr
if self._capacity_sufficient(size_to_add):
self._contents.append(object_)
self._recalc_capacity(size_to_add)
else:
self.__raise_for_exceeded_capacity()
def extend(
self,
collection: Collection[T],
) -> None:
size_to_add: int = len(collection)
if self.cap_property is not None:
if not all(hasattr(object_, self.cap_property) for object_ in collection):
raise AttributeError("Not all objects posses the wanted property")
if not all(
isinstance(getattr(object_, self.cap_property), int) for object_ in collection
):
raise TypeError("Capacity property must be an integer")
size_to_add = sum(getattr(object_, self.cap_property) for object_ in collection)
if self._capacity_sufficient(size_to_add):
self._contents.extend(collection)
self._recalc_capacity(size_to_add)
else:
self.__raise_for_exceeded_capacity()
def remove(
self,
object_: T,
) -> None:
size_to_add: int = 1
if self.cap_property is not None:
if not hasattr(object_, self.cap_property):
raise AttributeError("Object does not posses the wanted property")
attr = getattr(object_, self.cap_property)
if not isinstance(attr, int):
raise TypeError("Capacity property must be an integer")
size_to_add = attr
size_to_remove = (-1) * size_to_add
self._contents.remove(object_)
self._recalc_capacity(size_to_remove)