Coverage for src/extratools_core/seq/__init__.py: 53%

60 statements  

« prev     ^ index     » next       coverage.py v7.8.1, created at 2025-06-19 04:20 -0700

1import operator 

2from collections import Counter 

3from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence 

4from itertools import groupby, repeat 

5from typing import Any, cast 

6 

7from toolz import itertoolz, unique 

8 

9from ..typing import Comparable 

10from .common import iter_to_seq # noqa: F401 

11 

12 

13def sorted_by_rank[T]( 

14 data: Iterable[T], 

15 ranks: Iterable[Comparable], 

16 *, 

17 _reverse: bool = False, 

18) -> list[T]: 

19 return [ 

20 v 

21 for v, _ in sorted( 

22 zip(data, ranks, strict=True), 

23 key=lambda x: x[1], 

24 reverse=_reverse, 

25 ) 

26 ] 

27 

28 

29def compress[T]( 

30 data: Iterable[T], 

31 key: Callable[[T], Any] | None = None, 

32) -> Iterable[tuple[T, int]]: 

33 for k, g in groupby(data, key=key): 

34 yield (k, itertoolz.count(g)) 

35 

36 

37def decompress[T](data: Iterable[tuple[T, int]]) -> Iterable[T]: 

38 for k, n in data: 

39 yield from repeat(k, n) 

40 

41 

42def to_deltas[T]( 

43 data: Iterable[T], 

44 op: Callable[[T, T], T] = operator.sub, 

45) -> Iterable[T]: 

46 seq: Iterator[T] = iter(data) 

47 

48 curr: T | None = next(seq, None) 

49 if curr is None: 

50 return 

51 

52 yield curr 

53 

54 prev: T = curr 

55 for curr in seq: 

56 yield op(curr, prev) 

57 

58 prev = curr 

59 

60 

61def from_deltas[T]( 

62 data: Iterable[T], 

63 op: Callable[[T, T], T] = operator.add, 

64) -> Iterable[T]: 

65 seq: Iterator[T] = iter(data) 

66 

67 curr: T | None = next(seq, None) 

68 if curr is None: 

69 return 

70 

71 yield curr 

72 

73 prev: T = curr 

74 for curr in seq: 

75 res: T = op(prev, curr) 

76 yield res 

77 

78 prev = res 

79 

80 

81def key_frequencies[KT, VT]( 

82 *seqs: Iterable[KT], 

83 key: Callable[[KT], VT] | None = None, 

84) -> Mapping[VT, int]: 

85 c: Counter[VT] = Counter() 

86 for seq in seqs: 

87 c.update(cast("Iterable[VT]", unique(seq, key=key))) 

88 

89 return c 

90 

91 

92def add_until[T]( 

93 seq: Sequence[T], 

94 cond: Callable[[T], bool], 

95 *, 

96 op: Callable[[T, T], T] = operator.add, 

97 include_all: bool = True, 

98) -> Sequence[T]: 

99 i: int = 0 

100 

101 results: list[T] = [] 

102 

103 buffer: T | None = None 

104 while i < len(seq): 

105 element: T = seq[i] 

106 i += 1 

107 

108 if buffer is not None: 

109 element = op(buffer, element) 

110 

111 if cond(element): 

112 results.append(element) 

113 buffer = None 

114 else: 

115 buffer = element 

116 

117 if buffer is not None: 

118 if len(results) > 0: 

119 results[-1] = op(results[-1], buffer) 

120 else: 

121 results.append(buffer) 

122 

123 return results if include_all else list(filter(cond, results))