implement sized chunks, related to #6

2026-03-23 11:50:26 +01:00
parent 49d68cefa5
commit 320ff92c7d
2 changed files with 269 additions and 0 deletions
--- a/src/dopt_basics/batching.py
+++ b/src/dopt_basics/batching.py
@@ -0,0 +1,127 @@
+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)