| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Contains helpers that are applied to collections. |
| 4 | |
| 5 | This is similar in concept to the `collections` module. |
| 6 | """ |
| 7 | |
| 8 | import math |
| 9 | from collections import defaultdict |
| 10 | from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence |
| 11 | from typing import Generic, Literal, TypeVar |
| 12 | |
| 13 | from typing_extensions import TypeIs, assert_never, overload |
| 14 | |
| 15 | T = TypeVar("T") |
| 16 | |
| 17 | _K = TypeVar("_K", bound=Hashable) |
| 18 | _V = TypeVar("_V") |
| 19 | |
| 20 | |
| 21 | class LazyDict(Mapping[str, _V], Generic[_V]): |
| 22 | """Evaluates dictionary items only when they are accessed. |
| 23 | |
| 24 | Adapted from: https://stackoverflow.com/a/47212782/5082708 |
| 25 | """ |
| 26 | |
| 27 | def __init__(self, factory: dict[str, Callable[[], _V]]): |
| 28 | self._factory = factory |
| 29 | self._dict: dict[str, _V] = {} |
| 30 | |
| 31 | def __getitem__(self, key: str) -> _V: |
| 32 | if key not in self._dict: |
| 33 | if key not in self._factory: |
| 34 | raise KeyError(key) |
| 35 | self._dict[key] = self._factory[key]() |
| 36 | return self._dict[key] |
| 37 | |
| 38 | def __setitem__(self, key: str, value: Callable[[], _V]): |
| 39 | self._factory[key] = value |
| 40 | |
| 41 | def __iter__(self): |
| 42 | return iter(self._factory) |
| 43 | |
| 44 | def __len__(self): |
| 45 | return len(self._factory) |
| 46 | |
| 47 | |
| 48 | def as_list(maybe_list: Iterable[T]) -> list[T]: |
| 49 | """Convert iterable to list, unless it's already a list.""" |
| 50 | return maybe_list if isinstance(maybe_list, list) else list(maybe_list) |
| 51 | |
| 52 | |
| 53 | def is_list_of( |
| 54 | value: object, |
| 55 | typ: type[T] | tuple[type[T], ...], |
| 56 | *, |
| 57 | check: Literal["first", "all"] = "first", |
| 58 | ) -> TypeIs[list[T]]: |
| 59 | if not isinstance(value, list): |
| 60 | return False |
| 61 | |
| 62 | if check == "first": |
| 63 | return len(value) == 0 or isinstance(value[0], typ) |
| 64 | elif check == "all": |
| 65 | return all(isinstance(v, typ) for v in value) |
| 66 | |
| 67 | assert_never(check) |
| 68 | |
| 69 | |
| 70 | def is_list_of_numbers(value: object) -> bool: |
| 71 | """Check every item is an int or finite float, excluding booleans.""" |
| 72 | return isinstance(value, list) and all( |
| 73 | type(v) is int or (type(v) is float and math.isfinite(v)) for v in value |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | @overload |
| 78 | def common_prefix(items: Sequence[str]) -> str: ... |
| 79 | |
| 80 | |
| 81 | @overload |
| 82 | def common_prefix(items: Sequence[Sequence[T]]) -> Sequence[T]: ... |
| 83 | |
| 84 | |
| 85 | def common_prefix(items: Sequence[Sequence[T] | str]) -> Sequence[T] | str: |
| 86 | """Find the longest prefix common to all items.""" |
| 87 | if len(items) == 0: |
| 88 | return [] |
| 89 | if len(items) == 1: |
| 90 | return items[0] |
| 91 | |
| 92 | shortest = min(items, key=len) |
| 93 | if not shortest: |
| 94 | return shortest[:0] |
| 95 | |
| 96 | for match_len in range(1, len(shortest) + 1): |
| 97 | match = shortest[:match_len] |
| 98 | for item in items: |
| 99 | if item[:match_len] != match: |
| 100 | return shortest[: match_len - 1] |
| 101 | |
| 102 | return shortest |
| 103 | |
| 104 | |
| 105 | def chunk_list(lst: list[T], chunk_size: int) -> Generator[list[T]]: |
| 106 | """Yield successive chunk_size chunks from lst.""" |
| 107 | for i in range(0, len(lst), chunk_size): |
| 108 | yield lst[i : i + chunk_size] |
| 109 | |
| 110 | |
| 111 | def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]: |
| 112 | """Flatten a list of lists to a single list.""" |
| 113 | return [item for sublist in lists for item in sublist] |
| 114 | |
| 115 | |
| 116 | def full_groupby(values: Iterable[_V], *, key: Callable[[_V], _K]): |
| 117 | """Unlike [`itertools.groupby`][], groups are not broken by |
| 118 | non-contiguous data. |
| 119 | """ |
| 120 | groups = defaultdict[_K, list[_V]](list) |
| 121 | |
| 122 | for value in values: |
| 123 | groups[key(value)].append(value) |
| 124 | |
| 125 | return groups.items() |
| 126 | |
| 127 | |
| 128 | def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None: |
| 129 | """Swap values between two keys.""" |
| 130 | v1 = obj.get(key1) |
| 131 | v2 = obj.get(key2) |
| 132 | if v1 is not None: |
| 133 | obj[key2] = v1 |
| 134 | else: |
| 135 | obj.pop(key2, None) |
| 136 | if v2 is not None: |
| 137 | obj[key1] = v2 |
| 138 | else: |
| 139 | obj.pop(key1, None) |