Moving smart lights into smart_home to prepare for adding
[python_utils.git] / list_utils.py
index 182e2bc5c104908f39a15e4675021e6ed8a7c338..05512b564c2303372ff81660840bfd29dcba942b 100644 (file)
@@ -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]:
@@ -100,12 +127,75 @@ def dedup_list(lst: List[Any]) -> List[Any]:
 def uniq(lst: List[Any]) -> List[Any]:
     """
     Alias for dedup_list.
 def uniq(lst: List[Any]) -> List[Any]:
     """
     Alias for dedup_list.
-
     """
     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):
 def ngrams(lst: Sequence[Any], n):
+    """
+    Return the ngrams in the sequence.
+
+    >>> seq = 'encyclopedia'
+    >>> for _ in ngrams(seq, 3):
+    ...     _
+    'enc'
+    'ncy'
+    'cyc'
+    'ycl'
+    'clo'
+    'lop'
+    'ope'
+    'ped'
+    'edi'
+    'dia'
+
+    >>> seq = ['this', 'is', 'an', 'awesome', 'test']
+    >>> for _ in ngrams(seq, 3):
+    ...     _
+    ['this', 'is', 'an']
+    ['is', 'an', 'awesome']
+    ['an', 'awesome', 'test']
+    """
     for i in range(len(lst) - n + 1):
         yield lst[i:i + n]
 
     for i in range(len(lst) - n + 1):
         yield lst[i:i + n]