Coverage for src/extratools_core/strtools.py: 55%

1import gzip 

2from base64 import b64decode, b64encode 

3from collections.abc import Callable, Iterable 

4from typing import Literal 

5 

6import simple_zstd as zstd 

7 

8from .iter import iter_to_grams 

9from .seq.subseq import common_subseq, enumerate_subseqs 

10 

11 

12def str_to_grams( 

13 s: str, 

14 *, 

15 n: int, 

16 pad: str = '', 

17) -> Iterable[str]: 

18 if n < 1 or len(pad) > 1: 

19 raise ValueError 

20 

21 for c in iter_to_grams(s, n=n, pad=pad or None): 

22 yield ''.join(c) 

23 

24 

25def common_substr(a: str, b: str) -> str: 

26 return ''.join(common_subseq(a, b)) 

27 

28 

29def enumerate_substrs(s: str) -> Iterable[str]: 

30 return map(str, enumerate_subseqs(s)) 

31 

32 

33def compress( 

34 s: str, 

35 compress_func: Callable[[bytes], bytes] = gzip.compress, 

36) -> str: 

37 """ 

38 Compress a string by GZip + Base64 encoding. 

39 

40 https://base64.guru/developers/data-uri/gzip 

41 

42 Parameters 

43 ---------- 

44 s : str 

45 String to compress 

46 compress_func : Callable[[bytes], bytes] 

47 Optional function to compress (`gzip.compress` in default) 

48 

49 Returns 

50 ------- 

51 str 

52 Compressed string 

53 """ 

54 

55 return b64encode(compress_func(s.encode())).decode() 

56 

57 

58def encode( 

59 s: str, 

60 *, 

61 encoding: Literal["gzip", "zstd"] | None = None, 

62) -> str: 

63 match encoding: 

64 case "gzip": 

65 return compress(s, gzip.compress) 

66 case "zstd": 

67 return compress(s, zstd.compress) 

68 case None: 

69 return s 

70 case _: 

71 raise ValueError 

72 

73 

74def decompress( 

75 s: str, 

76 decompress_func: Callable[[bytes], bytes] = gzip.decompress, 

77) -> str: 

78 """ 

79 Decompress a string with GZip + Base64 encoding. 

80 

81 https://base64.guru/developers/data-uri/gzip 

82 

83 Parameters 

84 ---------- 

85 s : str 

86 String to decompress 

87 compress_func : Callable[[bytes], bytes] 

88 Optional function to decompress (`gzip.decompress` in default) 

89 

90 Returns 

91 ------- 

92 str 

93 Decompressed string 

94 """ 

95 

96 return decompress_func(b64decode(s.encode())).decode() 

97 

98 

99def decode( 

100 s: str, 

101 *, 

102 encoding: Literal["gzip", "zstd"] | None = None, 

103) -> str: 

104 match encoding: 

105 case "gzip": 

106 return decompress(s, gzip.decompress) 

107 case "zstd": 

108 return decompress(s, zstd.decompress) 

109 case None: 

110 return s 

111 case _: 

112 raise ValueError