lib.helpers

  1import os
  2from typing import Literal, TypeVar, TypeAlias, Generator
  3import sys
  4import random
  5import json
  6from math import ceil
  7from itertools import groupby
  8from copy import deepcopy
  9from itertools import islice
 10import loguru
 11from datetime import datetime
 12from icecream import ic
 13
 14T = TypeVar("T")
 15Strategy: TypeAlias = Literal["pick", "omit"]
 16"""Strategy for filtering items in a list."""
 17
 18
 19def setLogLevel(
 20    instance: "loguru.Logger",
 21    level: Literal["TRACE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "DEBUG",
 22):
 23    """Set the log level and format for a loguru.Logger instance.
 24
 25    Args:
 26        instance: The loguru.Logger instance to configure.
 27        level: The log level to set.
 28    """
 29    instance.remove()
 30    instance.add(
 31        sink=sys.stderr,
 32        level=level.upper(),
 33        format="<level>{level: ^8}</level> <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> <level>{message}</level>",
 34    )
 35
 36
 37# Filesystem
 38def readFile(filePath: str, mode: Literal["txt", "json"] = "txt"):
 39    """Read a file and return its contents.
 40
 41    Args:
 42        filePath: Path to the file.
 43        mode: `txt` to return raw str, `json` to return parsed JSON.
 44
 45    Returns:
 46        File contents as str or parsed JSON.
 47    """
 48    file = open(filePath, encoding="utf-8").read()
 49    # Parse JSON, otherwise return raw str
 50    return json.loads(file) if mode == "json" else file
 51
 52
 53def createFolder(path: str) -> str:
 54    """Create folder if non-existent, return its path.
 55
 56    Args:
 57        path: Path to the folder.
 58
 59    Returns:
 60        The path if created or exists, otherwise an empty string.
 61    """
 62    try:
 63        # ? Create parent directories recursively, ignore error if exists
 64        os.makedirs(path, exist_ok=True)
 65        return path
 66    except Exception as e:
 67        if isDirectory(path):
 68            return path
 69        else:
 70            loguru.logger.warning(
 71                "Error creating folder for path '{}' with exception {}", path, e
 72            )
 73            # Return empty string if failed => write to root
 74            return ""
 75
 76
 77def isFile(string: str) -> bool:
 78    """Returns True if the string is a path to a non-empty file."""
 79    return (
 80        isinstance(string, str)
 81        and os.path.isfile(string)
 82        and os.path.getsize(string) > 0
 83    )
 84
 85
 86def isDirectory(string: str) -> bool:
 87    """Returns True if the string is a directory path."""
 88    return os.path.isdir(string)
 89
 90
 91def coerceList(value, separate=False, strict=False) -> list:
 92    """Coerce a value to a list.
 93
 94    Args:
 95        value: The value to coerce.
 96        separate: If True, split str to characters.
 97            - (False) `"abc"` => `["abc"]`
 98            - (True) `"abc"` => `["a", "b", "c"]`
 99        strict: If True, force tuple to list.
100            - (False) `(a, b)` => `(a, b)`
101            - (True) `(a, b)` => `[(a, b)]`
102
103    Returns:
104        The coerced list.
105
106    Example:
107    - `1` or `[1]` => `[1]`
108    - `1, 2` => `[1, 2]`
109    """
110    passingTypes = list if strict else (list, tuple)
111    # Check if not None, can be []
112    if not value is None:
113        if isinstance(value, passingTypes):
114            return value
115        else:
116            return list(value) if separate else [value]
117
118
119def coerceNumericList(value, separator=" ") -> list[int]:
120    """Coerce a value to a list of integers.
121
122    Args:
123        value: The value to coerce.
124        separator: Separator for splitting strings.
125
126    Returns:
127        List of integers.
128
129    Example:
130    - (int) `5` => `[5]`
131    - (str) `"5 5"` => `[5, 5]`
132    """
133    if isinstance(value, (int, float)):
134        value = [value]
135    elif isinstance(value, str):
136        value = [int(side) for side in value.split(separator)]
137    return value
138
139
140def coerceTuple(value) -> tuple:
141    """Coerce a value to a tuple.
142
143    Args:
144        value: The value to coerce.
145
146    Returns:
147        The coerced tuple.
148
149    Example:
150    - `1` or `(1,)` => `(1,)`
151    - `1, 2` => `(1, 2)`
152    """
153    if isinstance(value, tuple):
154        return value
155    else:
156        return (value,) if not isinstance(value, (list, tuple)) else tuple(value)
157
158
159def expand(
160    input, n=2, defaultValue=None, format: Literal["auto", "list", "tuple"] = "auto"
161) -> list | tuple:
162    """Expand to iterable with length of `n`.
163
164    Args:
165        input: The input value(s) to expand.
166        n: Desired length.
167        defaultValue: Fallback value for cycling.
168        format: Output type.
169
170    Returns:
171        Iterable of length n.
172
173    Example:
174    - `n=3`: `a` => `[a, a, a]`
175    - `n=4`: `[1, 2]` => `[1, 2, 1, 2]`
176    """
177
178    def _getValueAtIndex(i: int):
179        try:
180            return input[i]
181        except:
182            return defaultValue or input[i % len(input)]
183
184    input = coerceList(input)
185    if format == "auto":
186        outputFormat = tuple if isinstance(input, tuple) else list
187    else:
188        outputFormat = tuple if format == "tuple" else list
189    return outputFormat([_getValueAtIndex(i) for i in range(n)])
190
191
192# Predicate
193def isInRange(start: int, end: int, value: int):
194    """Returns True if value is in the inclusive range [start, end]."""
195    return start <= value <= end
196
197
198# Strings
199def truncateString(string: str, length=10, suffix="...") -> str:
200    """
201    Truncate a string to a given length and append a suffix.
202
203    Example:
204        `(5)` `Hello there!`  => `Hello...`
205    """
206    return string[:length] + suffix
207
208
209# Lists
210def reverseList(input: list[T]) -> list[T]:
211    """Inverse values in iterable.
212
213    Returns:
214        The reversed list.
215    """
216    return input[::-1]
217
218
219def truncateList(input: list, length=3) -> str:
220    """
221    Truncate a list and describe remaining items.
222
223    Example:
224        (2) `[A, B, C, D]`  => `[A, B] and 2 more items`
225    """
226    remainder = len(input) - length
227    noun = "items" if remainder > 1 else "item"
228
229    output = f"{input[:length]}"
230
231    if remainder > 0:
232        output += f" and {remainder} more {noun}"
233
234    return output
235
236
237def inverseValues(values: list[int]) -> list:
238    """Negate all values in a list.
239
240    Example:
241       `[1, 2]` => `[-1, -2]`
242    """
243    return [-value for value in coerceList(values)]
244
245
246def flatten(items: list):
247    """Flatten a list of lists or tuples.
248
249    Returns:
250        Flattened list or original input if not all items are lists/tuples.
251    """
252    areItemsLists = all([isinstance(item, list) for item in items])
253    areItemsTuples = all([isinstance(item, tuple) for item in items])
254
255    if areItemsLists or areItemsTuples:
256        start = [] if areItemsLists else ()
257        return list(sum(items, start))
258    # Otherwise pass through unchanged
259    else:
260        isIterable = isinstance(items, (list, tuple))
261        if not isIterable:
262            loguru.logger.debug("Unable to flatten: {}", items)
263        return items
264
265
266def flattenTuples(input):
267    """Flatten nested tuples and lists into a single list of tuples.
268
269    Example:
270        `[(1, 2), [(3, 4), (5, 6)]]` => `[(1, 2), (3, 4), (5, 6)]`
271    """
272    output = []
273
274    if isinstance(input, tuple):
275        input = [input]
276
277    for item in input:
278        if isinstance(item, list):
279            for subitem in item:
280                output.append(subitem)
281        else:
282            output.append(item)
283    return output
284
285
286def intersect(lists: list[list], retainOrder=True) -> list:
287    """Return intersection of multiple lists.
288
289    Args:
290        lists: List of lists to intersect.
291        retainOrder: If True, order based on first list.
292
293    Returns:
294        List of intersected items.
295
296    Example:
297        `[a, b, c]` + `[b]` => `[b]`
298    """
299    result = set.intersection(*map(set, lists))
300    listA = lists[0]
301
302    return sorted(result, key=lambda x: listA.index(x)) if retainOrder else list(result)
303
304
305def zipWithMode(
306    groups: list[list], mode: Literal["drop", "fill"] = "fill"
307) -> list[list]:
308    """Zip lists with different modes.
309
310    Args:
311        groups: List of lists to zip.
312        mode: 'drop' to zip to shortest, 'fill' to zip to longest.
313
314    Returns:
315        Zipped list of lists.
316    """
317
318    def _fillGroup(group):
319        isShorter = len(group) < longest
320        return [group[i % len(group)] for i in range(longest)] if isShorter else group
321
322    if mode == "fill":
323        groups = [group if group else [group] for group in groups]  # None => [None]
324        longest = max([len(group) for group in groups])
325        groups = [_fillGroup(group) for group in groups]
326
327    return list(zip(*groups))
328
329
330def removeFalsy(
331    input: list, mode: Literal[None, False, "empty", "either"] = "either"
332) -> list:
333    """Remove falsy values from a list.
334
335    Args:
336        input: List to filter.
337        mode: 'None', 'False', 'empty', or 'either'.
338
339    Returns:
340        Filtered list.
341    """
342
343    def _checkItem(item):
344        isNotNone = item is not None
345        isNotFalse = item is not False
346        isNotEmpty = isinstance(item, list) and len(item)
347        isNotEither = isNotNone and isNotFalse and isNotEmpty
348        if mode is None:
349            return isNotNone
350        elif mode is False:
351            return isNotFalse
352        else:
353            return isNotEither
354
355    return list(filter(_checkItem, input))
356
357
358def removeNone(input: list) -> list:
359    """Remove `None` only from a list.
360
361    Returns:
362        List with None values removed.
363    """
364    return removeFalsy(input=input, mode=None)
365
366
367def removeFromList(input: list, remove: list) -> list:
368    """Remove specified items from a list.
369
370    Example:
371        `[a, b, c]` + `[b]` => `[a, c]`
372
373    Returns:
374        List with specified items removed.
375    """
376    output = []
377    remove = flatten(coerceList(remove))
378
379    for element in input:
380        if element not in remove:
381            output.append(element)
382
383    return output
384
385
386def dedupe(input: list, sortAlphabetically=False, caseSensitive=False, debug=False):
387    """Remove duplicates in list (= uniq).
388
389    Args:
390        input: List to deduplicate.
391        sortAlphabetically: If True, sort output.
392        caseSensitive: If False, dedupe case-insensitively.
393        debug: If True, log duplicates.
394
395    Returns:
396        Deduplicated list.
397    """
398    input = [item for item in input if item]  # Remove None
399    output = list()
400    uniq = set()
401
402    for item in input:
403        itemCase = (
404            item if caseSensitive or isinstance(item, (int, float)) else item.casefold()
405        )
406        if itemCase not in uniq:
407            uniq.add(itemCase)
408            output.append(item)
409        elif debug:
410            loguru.logger.trace("[Dedupe] {}", item)
411
412    if sortAlphabetically:
413        output = sorted(output)
414
415    return output
416
417
418def findClosestValue(values: list, value: int, sort=True, discardLarger=False) -> int:
419    """Return closest value in a sorted list.
420
421    If two values are equally close, return the smallest value.
422
423    Args:
424        values: List of values.
425        value: Value to compare.
426        sort: If True, sort the list.
427        discardLarger: If True, discard values larger than input.
428
429    Returns:
430        Closest value or None if list is empty.
431    """
432    from bisect import bisect_left
433
434    # Discard larger values
435    if discardLarger:
436        values = [v for v in values if v <= value]
437
438    if sort:
439        values = sorted(values)
440
441    pos = bisect_left(values, value)
442    if pos == 0:
443        return values[0] if len(values) else None
444    if pos == len(values):
445        return values[-1]
446    before = values[pos - 1]
447    after = values[pos]
448
449    if after - value < value - before:
450        return after
451    else:
452        return before
453
454
455def findClosestIndex(values: list, value: str) -> int:
456    """Find index of closest matching string in a list.
457
458    Useful for finding `sort()` key.
459
460    Example:
461        `["B", "F"], "Foo"` => returns `1`
462
463    Returns:
464        Index of closest match or None.
465    """
466    if not value:
467        return None
468
469    indexFound = None
470
471    while len(value) and not isinstance(indexFound, int):
472        try:
473            indexFound = values.index(value)
474        except:
475            pass
476
477        value = value[:-1]
478
479    return indexFound
480
481
482def isListOfOne(items: list) -> bool:
483    """Returns True if the list contains exactly one item.
484
485    Example:
486    - `[1, 2]` => `false`
487    - `1`      => `false`
488    - `[1]`    => `true`
489    """
490    return isinstance(items, list) and len(items) == 1
491
492
493def isFirst(
494    parent: list | int,
495    child: list | int,
496    mode: Literal["indices", "lists"] = "indices",
497):
498    """Returns True if child is the first element in parent.
499
500    Args:
501        parent: Iterable or count.
502        child: Item or index.
503        mode: 'indices' or 'lists'.
504    """
505    childIndex = (
506        child if isinstance(child, int) and mode == "indices" else parent.index(child)
507    )
508    return childIndex == 0
509
510
511def isLast(
512    parent: list | int,
513    child: list | int,
514    mode: Literal["indices", "lists"] = "indices",
515):
516    """Returns True if child is the last element in parent.
517
518    Args:
519        parent: Iterable or count.
520        child: Item or index.
521        mode: 'indices' or 'lists'.
522    """
523    childIndex = (
524        child if isinstance(child, int) and mode == "indices" else parent.index(child)
525    )
526    parentCount = len(parent) if not isinstance(parent, int) else parent
527    return childIndex == parentCount - 1
528
529
530def isNth(index: int, nth: int = 3):
531    """Returns True if index is every nth element.
532
533    Args:
534        index: Index to check.
535        nth: Nth interval.
536
537    Example:
538        `nth=3` => Every third
539    """
540    return (index + 1) % nth == 0
541
542
543def pickFirst(value):
544    """Get first value from iterable if iterable.
545
546    Returns:
547        First value or value itself if not iterable.
548    """
549    if isinstance(value, (list, tuple)):
550        return value[0]
551    else:
552        return value
553
554
555def pickLast(value):
556    """Get last value from iterable if iterable.
557
558    Returns:
559        Last value or value itself if not iterable.
560    """
561    if isinstance(value, (list, tuple)):
562        return value[-1]
563    else:
564        return value
565
566
567def pickRandom(items: list) -> tuple:
568    """Pick 1 item at random and return remaining items.
569
570    Returns:
571        Tuple of (item, remainingItems).
572    """
573    item = random.choice(items)
574    itemIndex = items.index(item)
575    items.pop(itemIndex)
576
577    return item, items
578
579
580def pickExtremes(values: list):
581    """Get first and last from list.
582
583    Example:
584        `[1, 2, 3, 4]` => `(1, 4)`
585
586    Returns:
587        Tuple of (first, last) or values itself if length <= 1.
588    """
589    if len(values) > 1:
590        return values[0], values[-1]
591    else:
592        return values
593
594
595def pushToEnd(items: list[T], index=0):
596    """Move item at index to the end and return modified list.
597
598    Example:
599        `ABCD` `index=0` => `BCDA`
600
601    Returns:
602        Modified list.
603    """
604    result = items.copy()  # Avoid mutating original
605    result.append(result.pop(index))
606    return result
607
608
609def _splitToGroups(input: list | str, groups=2) -> Generator:
610    """Split input into groups of approximately equal size.
611
612    Args:
613        input: List or string to split.
614        groups: Number of groups.
615
616    Yields:
617        Groups as tuples or strings.
618    """
619    it = iter(input)
620    while batch := tuple(islice(it, ceil(len(input) / groups))):
621        yield "".join(batch) if isinstance(input, str) else batch
622
623
624def halven(input: list | str):
625    """Split input into two halves.
626    - `even`: two equal halves
627    - `odd`: larger/smaller half
628
629    Returns:
630        List of two halves (even or odd sized).
631    """
632    return list(_splitToGroups(input, 2))
633
634
635def splitEvenly(input: list | str, groups=3):
636    """Split input into evenly sized groups.
637
638    Args:
639        input: List or string to split.
640        groups: Number of groups.
641
642    Example:
643        `groups=2`, `ABCD` => `AB CD`
644
645    Returns:
646        List of groups.
647    """
648    return list(_splitToGroups(input, groups))
649
650
651def toChunks(items: list[T], groups=2) -> list[list[T]]:
652    """Combine groups in alternating manner.
653
654    Example:
655        `groups=2`, `ABCD` => `AC BD`
656
657    Returns:
658        List of chunked groups.
659    """
660    return [items[start::groups] for start in range(groups)]
661
662
663def shuffleAtRandomSegment(items: list[T]) -> list[T]:
664    """Shuffle list at a random segment to maintain order.
665
666    Returns:
667        List shuffled at a random index.
668    """
669    randomItem = random.choice(items)
670    randomIndex = items.index(randomItem)
671    return flatten([items[randomIndex:], items[:randomIndex]])
672
673
674def shuffleAtInterval(items: list[T], n: int = 2) -> list[T]:
675    """
676    Rearrange a list by grouping every nth item together, then flattening the result.
677
678    Example:
679    - For items `[A, B, C, D, E, F]` and n `3`:
680    - The list is split into 3 groups: `[A, D], [B, E], [C, F]`
681    - The result is: `[A, D, B, E, C, F]`
682
683    Returns:
684        A new list with items reordered by interval.
685    """
686    return sum(toChunks(items, n), [])
687
688
689def spaceEvenly(items: list, size=2):
690    """Get n items from list, evenly spaced.
691
692    Args:
693        items: List of items.
694        size: Number of items to get.
695
696    Example:
697    - 1 => get midpoint
698    - 2 => get first and last
699
700    Returns:
701        List of evenly spaced items.
702    """
703
704    def _getSpacedIndices(start, stop, length):
705        return [round(start + x * (stop - start) / (length - 1)) for x in range(length)]
706
707    itemCount = len(items)
708    if size == 1:
709        # Return midpoint
710        indices = [len(items) >> 1]
711    else:
712        # Never return size larger than itemCount
713        indices = _getSpacedIndices(0, itemCount - 1, min(itemCount, size))
714
715    return [items[i] for i in indices]
716
717
718# Object: dict
719def pick(object: dict, keys: list[str]) -> dict:
720    """Pick properties by `key`.
721
722    Args:
723        object: Dictionary to pick from.
724        keys: Keys to pick.
725
726    Returns:
727        Dictionary with only picked keys.
728    """
729    return {k: v for k, v in object.items() if k in coerceList(keys)}
730
731
732def omit(object: dict, keys: list[str]) -> dict:
733    """Remove properties by `key`.
734
735    Args:
736        object: Dictionary to omit from.
737        keys: Keys to omit.
738
739    Returns:
740        Dictionary with omitted keys.
741    """
742    return {k: v for k, v in object.items() if k not in coerceList(keys)}
743
744
745def omitBy(object: dict, omit=[0, None, []]) -> dict:
746    """Remove properties by `value`.
747
748    Remove falsy by default: `0`, `None`, `[]`.
749
750    Args:
751        object: Dictionary to filter.
752        omit: Values to omit.
753
754    Returns:
755        Dictionary with omitted values.
756    """
757    return {k: v for k, v in object.items() if v not in coerceList(omit)}
758
759
760# Collections: list[dict]
761def merge(dicts: list[dict]):
762    """Merge multiple instances of `dict`.
763
764    What happens to values:
765    - `str` gets overwritten by last occurrence
766    - `list` gets extended (`flatten` + `dedupe`)
767
768    Args:
769        dicts: List of dictionaries to merge.
770
771    Returns:
772        Merged dictionary.
773    """
774    result = dict()
775
776    for d in dicts:
777        for [key, value] in d.items():
778            hasKey = result.get(key)
779            if hasKey:
780                if isinstance(value, list):
781                    result[key] = dedupe(flatten([result[key], value]))
782                else:
783                    result[key] = value
784            else:
785                result[key] = value
786
787    return result
788
789
790def findByKey(items: list[dict], key: str, value, defaultValue=None) -> dict | None:
791    """Find item in list of dictionaries/objects by key.
792
793    Args:
794        items: List of dictionaries or objects.
795        key: Key to search for.
796        value: Value to match.
797        defaultValue: Value to return if not found.
798
799    Returns:
800        The found item or defaultValue.
801    """
802
803    def getValue(item, key):
804        # Use built-in method for dictionaries
805        if isinstance(item, dict):
806            return item.get(key)
807        # Otherwise assume it’s an object
808        else:
809            return getattr(item, key)
810
811    return next((item for item in items if getValue(item, key) == value), defaultValue)
812
813
814def sortByKey(items: list[dict], key: str, defaultValue: int | str = None):
815    """Sort list of dictionaries by key.
816
817    Args:
818        items: List of dictionaries.
819        key: Key to sort by.
820        defaultValue: Value to use if key is missing.
821
822    Returns:
823        Sorted list or original if sorting fails.
824    """
825    # May be []
826    try:
827        return sorted(items, key=lambda item: item.get(key, defaultValue))
828    except:
829        return items
830
831
832def groupByKey(items: list[dict], key: str) -> list[list[dict]]:
833    """Group list of dictionaries by `key`.
834
835    Args:
836        items: List of dictionaries.
837        key: Key to group by.
838
839    Returns:
840        Grouped items as an iterator.
841    """
842    return groupby(items, key=lambda item: item.get(key))
843
844
845def reject(items: list[dict], key: str, value):
846    """Remove items matching `key`:`value`.
847
848    Args:
849        items: List of dictionaries.
850        key: Key to check.
851        value: Value to reject.
852
853    Returns:
854        Filtered list.
855    """
856    return [item for item in items if item.get(key) != value]
857
858
859def mergeCollections(
860    main: list[dict], additions: list[dict], key: str = "id", caseSensitive=True
861) -> list[dict]:
862    """Merge two instances of `list[dict]` by key.
863
864    Args:
865        main: Main list of dictionaries.
866        additions: List of dictionaries to merge in.
867        key: Key to match for merging.
868        caseSensitive: If False, dedupe case-insensitively.
869
870    Returns:
871        Merged list of dictionaries.
872    """
873    # Avoid modifying original collection
874    merged = deepcopy(main)
875    for addition in additions:
876        foundMain = findByKey(merged, key, addition.get(key))
877
878        # Update existing dictionary
879        if foundMain:
880            for property in addition:
881                valueMain = foundMain.get(property)
882                valueAdd = addition.get(property)
883                # Update value if new
884                if not valueMain == valueAdd:
885                    # Concat if both are lists
886                    if all([isinstance(v, list) for v in [valueMain, valueAdd]]):
887                        valueNew = dedupe(
888                            flatten([*valueMain, *valueAdd]),
889                            caseSensitive=caseSensitive,
890                        )
891                    # Otherwise overwrite
892                    else:
893                        valueNew = valueAdd
894                    foundMain.update({property: valueNew})
895        # Otherwise insert as last
896        else:
897            merged.append(addition)
898    return merged
899
900
901def flatMap(collection: list[dict], key: str) -> list:
902    """Map all `key` values into a single list.
903
904    Args:
905        collection: List of dictionaries.
906        key: Key to extract.
907
908    Returns:
909        Flattened list of all key values.
910    """
911    # Return empty [] if key not found to facilitate flatten
912    return flatten([item.get(key, []) for item in collection])
913
914
915# Misc
916def getTimestamp(format: Literal["machine", "human"]) -> str:
917    """Get current timestamp in specified format.
918
919    Example:
920        - `machine` => `2024-06-14`
921        - `human`   => `Jun 14, 2024`
922    """
923    now = datetime.now()
924    if format == "machine":
925        return now.strftime("%Y-%m-%d")
926    else:
927        return now.strftime("%b %d, %Y")