A class for reordering and batching elements of an array. This class allows for sorting an array based on a provided sorting function, grouping elements based on a grouping function, and generating batches from the sorted and grouped data.
| 353 | |
| 354 | |
| 355 | class Collator: |
| 356 | """ |
| 357 | A class for reordering and batching elements of an array. |
| 358 | |
| 359 | This class allows for sorting an array based on a provided sorting function, grouping elements based on a grouping function, and generating batches from the sorted and grouped data. |
| 360 | """ |
| 361 | |
| 362 | def __init__( |
| 363 | self, |
| 364 | arr: List, |
| 365 | sort_fn: Callable, |
| 366 | group_fn: Callable = lambda x: x[1], |
| 367 | grouping: bool = False, |
| 368 | ) -> None: |
| 369 | self.grouping = grouping |
| 370 | self.fn = sort_fn |
| 371 | self.group_fn = lambda x: group_fn(x[1]) # first index are enumerated indices |
| 372 | self.reorder_indices: List = [] |
| 373 | self.size = len(arr) |
| 374 | self.arr_with_indices: Iterable[Any] = tuple(enumerate(arr)) # [indices, (arr)] |
| 375 | if self.grouping is True: |
| 376 | self.group_by_index() |
| 377 | |
| 378 | def group_by_index(self) -> None: |
| 379 | self.arr_with_indices = self.group( |
| 380 | self.arr_with_indices, fn=self.group_fn, values=False |
| 381 | ) |
| 382 | |
| 383 | def get_batched(self, n: int = 1, batch_fn: Optional[Callable] = None) -> Iterator: |
| 384 | """ |
| 385 | Generates and yields batches from the reordered array. |
| 386 | |
| 387 | Parameters: |
| 388 | - n (int): The size of each batch. Defaults to 1. |
| 389 | - batch_fn (Optional[Callable[[int, Iterable], int]]): A function to determine the size of each batch. Defaults to None. |
| 390 | |
| 391 | Yields: |
| 392 | Iterator: An iterator over batches of reordered elements. |
| 393 | """ |
| 394 | if self.grouping: |
| 395 | for ( |
| 396 | key, |
| 397 | values, |
| 398 | ) in self.arr_with_indices.items(): # type: ignore |
| 399 | values = self._reorder(values) |
| 400 | batch = self.get_chunks(values, n=n, fn=batch_fn) |
| 401 | yield from batch |
| 402 | else: |
| 403 | values = self._reorder(self.arr_with_indices) # type: ignore |
| 404 | batch = self.get_chunks(values, n=n, fn=batch_fn) |
| 405 | yield from batch |
| 406 | |
| 407 | def _reorder(self, arr: Union[List, Tuple[Tuple[int, Any], ...]]) -> List: |
| 408 | """ |
| 409 | Reorders the elements in the array based on the sorting function. |
| 410 | |
| 411 | Parameters: |
| 412 | - arr (Union[List, Tuple[Tuple[int, Any], ...]]): The array or iterable to be reordered. |
no outgoing calls
no test coverage detected