import collections from collections.abc import Iterable, Sequence from typing import Any, Callable, Generator, Literal, TypeVar, overload _T = TypeVar('_T') _R = TypeVar('_R') @overload def partition(lst: Iterable[_T], f: Callable[[_T], _R], save_keys: Literal[True]) -> collections.OrderedDict[_R, list[_T]]: ... @overload def partition(lst: Iterable[_T], f: Callable[[_T], _R], save_keys: Literal[False]) -> list[list[_T]]: ... def partition(lst: Iterable[_T], f: Callable[[_T], _R], save_keys: bool = False) -> list[list[_T]] | collections.OrderedDict[_R, list[_T]]: """partition(lst, f, save_keys = False) -> list Partitions an iterable into sublists using a function to specify which group they belong to. It works by calling `f` on every element and saving the results into an :class:`collections.OrderedDict`. Arguments: lst: The iterable to partition f(function): The function to use as the partitioner. save_keys(bool): Set this to True, if you want the OrderedDict returned instead of just the values Example: >>> partition([1,2,3,4,5], lambda x: x&1) [[1, 3, 5], [2, 4]] >>> partition([1,2,3,4,5], lambda x: x%3, save_keys=True) == collections.OrderedDict([(1, [1, 4]), (2, [2, 5]), (0, [3])]) True """ d: collections.OrderedDict[_R, list[_T]] = collections.OrderedDict() for l in lst: c = f(l) s = d.setdefault(c, []) s.append(l) if save_keys: return d else: return list(d.values()) @overload def group(n: int, lst: list[_T], underfull_action: Literal['ignore', 'drop', 'fill'] = 'ignore', fill_value: _T | None = None) -> list[list[_T]]: ... @overload def group(n: int, lst: tuple[_T, ...], underfull_action: Literal['ignore', 'drop', 'fill'] = 'ignore', fill_value: _T | None = None) -> list[tuple[_T, ...]]: ... @overload def group(n: int, lst: str, underfull_action: Literal['ignore', 'drop', 'fill'] = 'ignore', fill_value: str | None = None) -> list[str]: ... @overload def group(n: int, lst: bytes, underfull_action: Literal['ignore', 'drop', 'fill'] = 'ignore', fill_value: bytes | None = None) -> list[bytes]: ... def group(n: int, lst: Sequence[_T], underfull_action: Literal['ignore', 'drop', 'fill'] = 'ignore', fill_value: _T | None = None) -> list[Any]: """group(n, lst, underfull_action = 'ignore', fill_value = None) -> list Split sequence into subsequences of given size. If the values cannot be evenly distributed among into groups, then the last group will either be returned as is, thrown out or padded with the value specified in fill_value. Arguments: n (int): The size of resulting groups lst: The list, tuple or string to group underfull_action (str): The action to take in case of an underfull group at the end. Possible values are 'ignore', 'drop' or 'fill'. fill_value: The value to fill into an underfull remaining group. Returns: A list containing the grouped values. Example: >>> group(3, "ABCDEFG") ['ABC', 'DEF', 'G'] >>> group(3, 'ABCDEFG', 'drop') ['ABC', 'DEF'] >>> group(3, 'ABCDEFG', 'fill', 'Z') ['ABC', 'DEF', 'GZZ'] >>> group(3, list('ABCDEFG'), 'fill') [['A', 'B', 'C'], ['D', 'E', 'F'], ['G', None, None]] >>> group(2, tuple('1234'), 'fill') [('1', '2'), ('3', '4')] """ if underfull_action not in ['ignore', 'drop', 'fill']: raise ValueError("group(): underfull_action must be either 'ignore', 'drop' or 'fill'") fill: Any = fill_value if underfull_action == 'fill': if isinstance(lst, tuple): fill = (fill_value,) elif isinstance(lst, list): fill = [fill_value] elif isinstance(lst, (bytes, str)): if not isinstance(fill_value, (bytes, str)): raise ValueError("group(): cannot fill a string with a non-string") else: raise ValueError("group(): 'lst' must be either a tuple, list or string") out: list[Any] = [] for i in range(0, len(lst), n): chunk: Any = lst[i:i+n] if len(chunk) < n: if underfull_action == 'ignore': pass elif underfull_action == 'drop': continue else: chunk = chunk + fill * (n - len(chunk)) out.append(chunk) return out def concat(l: list[list[Any]]) -> list[Any]: """concat(l) -> list Concats a list of lists into a list. Example: >>> concat([[1, 2], [3]]) [1, 2, 3] """ res = [] for k in l: res.extend(k) return res def concat_all(*args: Any) -> list[Any]: """concat_all(*args) -> list Concats all the arguments together. Example: >>> concat_all(0, [1, (2, 3)], [([[4, 5, 6]])]) [0, 1, 2, 3, 4, 5, 6] """ def go(arg: Any, output: list[Any]) -> list[Any]: if isinstance(arg, (tuple, list)): for e in arg: go(e, output) else: output.append(arg) return output return go(args, []) def ordlist(s: str) -> list[int]: """ordlist(s) -> list Turns a string into a list of the corresponding ascii values. Example: >>> ordlist("hello") [104, 101, 108, 108, 111] """ return list(map(ord, s)) def unordlist(cs: list[int]) -> str: """unordlist(cs) -> str Takes a list of ascii values and returns the corresponding string. Example: >>> unordlist([104, 101, 108, 108, 111]) 'hello' """ return ''.join(chr(c) for c in cs) def findall(haystack: Sequence[_T], needle: _T | Sequence[_T]) -> Generator[int, None, None]: """findall(l, e) -> l Generate all indices of needle in haystack, using the Knuth-Morris-Pratt algorithm. Example: >>> foo = findall([1,2,3,4,4,3,4,2,1], 4) >>> next(foo) 3 >>> next(foo) 4 >>> next(foo) 6 >>> list(foo) # no more appearances [] >>> list(findall("aaabaaabc", "aab")) [1, 5] """ def __kmp_table(W: Sequence[_T]) -> list[int]: pos = 1 cnd = 0 T: list[int] = [] T.append(-1) T.append(0) while pos < len(W): if W[pos] == W[cnd]: cnd += 1 pos += 1 T.append(cnd) elif cnd > 0: cnd = T[cnd] else: pos += 1 T.append(0) return T def __kmp_search(S: Sequence[_T], W: Sequence[_T]) -> Generator[int, None, None]: m = 0 i = 0 T = __kmp_table(W) while m + i < len(S): if S[m + i] == W[i]: i += 1 if i == len(W): yield m m += i - T[i] i = max(T[i], 0) else: m += i - T[i] i = max(T[i], 0) def __single_search(S: Sequence[_T], w: _T) -> Generator[int, None, None]: for i, v in enumerate(S): if v == w: yield i if not isinstance(needle, Sequence): needle = [needle] if len(needle) == 1: return __single_search(haystack, needle[0]) else: return __kmp_search(haystack, needle)