pwntools/pwnlib/util/lists.py

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)