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