Coverage for src/extratools_core/itertools.py: 98%

57 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-04-07 08:10 -0700

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

2from itertools import chain, count, repeat 

3from typing import cast 

4 

5from toolz.itertoolz import sliding_window 

6 

7from .dict import invert 

8from .seq import iter_to_seq 

9from .typing import Comparable 

10 

11 

12def iter_to_grams[T]( 

13 data: Iterable[T], 

14 *, 

15 n: int, 

16 pad: T | None = None, 

17) -> Iterable[Sequence[T]]: 

18 if pad is not None: 

19 data = chain( 

20 repeat(pad, n - 1), 

21 data, 

22 repeat(pad, n - 1), 

23 ) 

24 

25 return sliding_window(n, data) 

26 

27 

28def is_sorted[T]( 

29 data: Iterable[T], 

30 *, 

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

32 reverse: bool = False, 

33) -> bool: 

34 local_key: Callable[[T], Comparable] 

35 if key is None: 

36 def default_key(v: T) -> Comparable: 

37 return cast("Comparable", v) 

38 

39 local_key = default_key 

40 else: 

41 local_key = key 

42 

43 return all( 

44 ( 

45 local_key(prev) >= local_key(curr) if reverse 

46 else local_key(prev) <= local_key(curr) 

47 ) 

48 for prev, curr in sliding_window(2, data) 

49 ) 

50 

51 

52def filter_by_positions[T](poss: Iterable[int], data: Iterable[T]) -> Iterable[T]: 

53 p: Iterator[int] = iter(poss) 

54 

55 pos: int | None = next(p, None) 

56 if pos is None: 

57 return 

58 

59 for i, v in enumerate(data): 

60 if i == pos: 

61 yield v 

62 

63 pos = next(p, None) 

64 if pos is None: 

65 return 

66 

67 

68def filter_by_others[T](func: Callable[[T, T], bool], data: Iterable[T]) -> Iterable[T]: 

69 seq: Sequence[T] = iter_to_seq(data) 

70 

71 filtered_ids: set[int] = set(range(len(seq))) 

72 

73 for i, x in enumerate(seq): 

74 remove: bool = False 

75 for j in filtered_ids: 

76 if i == j: 

77 continue 

78 

79 if not func(x, seq[j]): 

80 remove = True 

81 break 

82 

83 if remove: 

84 filtered_ids.remove(i) 

85 

86 for i in filtered_ids: 

87 yield seq[i] 

88 

89 

90def remap[KT, VT]( 

91 data: Iterable[KT], 

92 mapping: Mapping[KT, VT], 

93 *, 

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

95) -> Iterable[VT]: 

96 local_key: Callable[[KT], VT] 

97 if key is None: 

98 inverted_mapping: Mapping[VT, KT] = invert(mapping) 

99 c = count(start=0) 

100 

101 def default_key(_: KT) -> VT: 

102 while True: 

103 v: int = next(c) 

104 if v not in inverted_mapping: 

105 return cast("VT", v) 

106 

107 local_key = default_key 

108 else: 

109 local_key = key 

110 

111 k: KT 

112 for k in data: 

113 yield mapping.setdefault(k, local_key(k))