Optionally surface exceptions that happen under executors by reading

[python_utils.git] / list_utils.py

diff --git a/list_utils.py b/list_utils.py
index a8030e30812ceb804a6c45ff386a7ab159fa2911..b65ab0d69fab36196a94b4ef3327fae63b88f48d 100644 (file)
--- a/list_utils.py
+++ b/list_utils.py
@@ -2,7 +2,7 @@
 
 from collections import Counter
 from itertools import islice
 
 from collections import Counter
 from itertools import islice

-from typing import Any, Iterator, List, Mapping, Sequence
+from typing import Any, Iterator, List, Mapping, Sequence, Tuple

 
 
 def shard(lst: List[Any], size: int) -> Iterator[Any]:
 
 
 def shard(lst: List[Any], size: int) -> Iterator[Any]:


@@ -48,6 +48,23 @@ def prepend(item: Any, lst: List[Any]) -> List[Any]:
     return lst
 
 
     return lst
 
 

+def remove_list_if_one_element(lst: List[Any]) -> Any:
+    """
+    Remove the list and return the 0th element iff its length is one.
+
+    >>> remove_list_if_one_element([1234])
+    1234
+
+    >>> remove_list_if_one_element([1, 2, 3, 4])
+    [1, 2, 3, 4]
+
+    """
+    if len(lst) == 1:
+        return lst[0]
+    else:
+        return lst
+
+

 def population_counts(lst: List[Any]) -> Mapping[Any, int]:
     """
     Return a population count mapping for the list (i.e. the keys are
 def population_counts(lst: List[Any]) -> Mapping[Any, int]:
     """
     Return a population count mapping for the list (i.e. the keys are


@@ -61,29 +78,39 @@ def population_counts(lst: List[Any]) -> Mapping[Any, int]:
     return Counter(lst)
 
 
     return Counter(lst)
 
 

-def most_common_item(lst: List[Any]) -> Any:
+def most_common(lst: List[Any], *, count=1) -> Any:

 
     """
     Return the most common item in the list.  In the case of ties,
     which most common item is returned will be random.
 
 
     """
     Return the most common item in the list.  In the case of ties,
     which most common item is returned will be random.
 

-    >>> most_common_item([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4])
+    >>> most_common([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4])

     3
 
     3
 

+    >>> most_common([1, 1, 1, 2, 2, 3, 3, 3, 3, 4, 4], count=2)
+    [3, 1]
+

     """
     """

-    return population_counts(lst).most_common(1)[0][0]
+    p = population_counts(lst)
+    return remove_list_if_one_element([_[0] for _ in p.most_common()[0:count]])

 
 
 
 

-def least_common_item(lst: List[Any]) -> Any:
+def least_common(lst: List[Any], *, count=1) -> Any:

     """
     Return the least common item in the list.  In the case of
     ties, which least common item is returned will be random.
 
     """
     Return the least common item in the list.  In the case of
     ties, which least common item is returned will be random.
 

-    >>> least_common_item([1, 1, 1, 2, 2, 3, 3, 3, 4])
+    >>> least_common([1, 1, 1, 2, 2, 3, 3, 3, 4])

     4
 
     4
 

+    >>> least_common([1, 1, 1, 2, 2, 3, 3, 3, 4], count=2)
+    [4, 2]
+

     """
     """

-    return population_counts(lst).most_common()[-1][0]
+    p = population_counts(lst)
+    mc = p.most_common()[-count:]
+    mc.reverse()
+    return remove_list_if_one_element([_[0] for _ in mc])

 
 
 def dedup_list(lst: List[Any]) -> List[Any]:
 
 
 def dedup_list(lst: List[Any]) -> List[Any]:


@@ -104,6 +131,46 @@ def uniq(lst: List[Any]) -> List[Any]:
     return dedup_list(lst)
 
 
     return dedup_list(lst)
 
 

+def contains_duplicates(lst: List[Any]) -> bool:
+    """
+    Does the list contian duplicate elements or not?
+
+    >>> lst = [1, 2, 1, 3, 3, 4, 4, 5, 6, 1, 3, 4]
+    >>> contains_duplicates(lst)
+    True
+
+    >>> contains_duplicates(dedup_list(lst))
+    False
+
+    """
+    seen = set()
+    for _ in lst:
+        if _ in seen:
+            return True
+        seen.add(_)
+    return False
+
+
+def all_unique(lst: List[Any]) -> bool:
+    """
+    Inverted alias for contains_duplicates.
+    """
+    return not contains_duplicates(lst)
+
+
+def transpose(lst: List[Any]) -> List[Any]:
+    """
+    Transpose a list of lists.
+
+    >>> lst = [[1, 2], [3, 4], [5, 6]]
+    >>> transpose(lst)
+    [[1, 3, 5], [2, 4, 6]]
+
+    """
+    transposed = zip(*lst)
+    return [list(_) for _ in transposed]
+
+

 def ngrams(lst: Sequence[Any], n):
     """
     Return the ngrams in the sequence.
 def ngrams(lst: Sequence[Any], n):
     """
     Return the ngrams in the sequence.


@@ -130,9 +197,91 @@ def ngrams(lst: Sequence[Any], n):
     ['an', 'awesome', 'test']
     """
     for i in range(len(lst) - n + 1):
     ['an', 'awesome', 'test']
     """
     for i in range(len(lst) - n + 1):

-        yield lst[i:i + n]
+        yield lst[i : i + n]
+
+
+def permute(seq: Sequence[Any]):
+    """
+    Returns all permutations of a sequence; takes O(N^2) time.
+
+    >>> for x in permute('cat'):
+    ...     print(x)
+    cat
+    cta
+    act
+    atc
+    tca
+    tac
+
+    """
+    yield from _permute(seq, "")
+
+
+def _permute(seq: Sequence[Any], path):
+    if len(seq) == 0:
+        yield path
+
+    for i in range(len(seq)):
+        car = seq[i]
+        left = seq[0:i]
+        right = seq[i + 1 :]
+        cdr = left + right
+        yield from _permute(cdr, path + car)
+
+
+def binary_search(
+    lst: Sequence[Any], target: Any, *, sanity_check=False
+) -> Tuple[bool, int]:
+    """Performs a binary search on lst (which must already be sorted).
+    Returns a Tuple composed of a bool which indicates whether the
+    target was found and an int which indicates the index closest to
+    target whether it was found or not.
+
+    >>> a = [1, 4, 5, 6, 7, 9, 10, 11]
+    >>> binary_search(a, 4)
+    (True, 1)
+
+    >>> binary_search(a, 12)
+    (False, 8)
+
+    >>> binary_search(a, 3)
+    (False, 1)
+
+    >>> binary_search(a, 2)
+    (False, 1)
+
+    >>> a.append(9)
+    >>> binary_search(a, 4, sanity_check=True)
+    Traceback (most recent call last):
+    ...
+    AssertionError
+
+    """
+    if sanity_check:
+        last = None
+        for x in lst:
+            if last is not None:
+                assert x >= last  # This asserts iff the list isn't sorted
+            last = x  # in ascending order.
+    return _binary_search(lst, target, 0, len(lst) - 1)
+
+
+def _binary_search(
+    lst: Sequence[Any], target: Any, low: int, high: int
+) -> Tuple[bool, int]:
+    if high >= low:
+        mid = (high + low) // 2
+        if lst[mid] == target:
+            return (True, mid)
+        elif lst[mid] > target:
+            return _binary_search(lst, target, low, mid - 1)
+        else:
+            return _binary_search(lst, target, mid + 1, high)
+    else:
+        return (False, low)

 
 
 if __name__ == '__main__':
     import doctest
 
 
 if __name__ == '__main__':
     import doctest

+

     doctest.testmod()
     doctest.testmod()