codekingpro/portable-devtools

1# Copyright 2022-present, the HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""Contains a utility to iterate by chunks over an iterator."""15 16import itertools17from collections.abc import Iterable18from typing import TypeVar19 20 21T = TypeVar("T")22 23 24def chunk_iterable(iterable: Iterable[T], chunk_size: int) -> Iterable[Iterable[T]]:25    """Iterates over an iterator chunk by chunk.26 27    Taken from https://stackoverflow.com/a/8998040.28    See also https://github.com/huggingface/huggingface_hub/pull/920#discussion_r938793088.29 30    Args:31        iterable (`Iterable`):32            The iterable on which we want to iterate.33        chunk_size (`int`):34            Size of the chunks. Must be a strictly positive integer (e.g. >0).35 36    Example:37 38    ```python39    >>> from huggingface_hub.utils import chunk_iterable40 41    >>> for items in chunk_iterable(range(17), chunk_size=8):42    ...     print(items)43    # [0, 1, 2, 3, 4, 5, 6, 7]44    # [8, 9, 10, 11, 12, 13, 14, 15]45    # [16] # smaller last chunk46    ```47 48    Raises:49        [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)50            If `chunk_size` <= 0.51 52    > [!WARNING]53    > The last chunk can be smaller than `chunk_size`.54    """55    if not isinstance(chunk_size, int) or chunk_size <= 0:56        raise ValueError("`chunk_size` must be a strictly positive integer (>0).")57 58    iterator = iter(iterable)59    while True:60        try:61            next_item = next(iterator)62        except StopIteration:63            return64        yield itertools.chain((next_item,), itertools.islice(iterator, chunk_size - 1))65