Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/compression_utils.py: 75%

1import asyncio 

2import zlib 

3from concurrent.futures import Executor 

4from typing import Optional, cast 

5 

6try: 

7 try: 

8 import brotlicffi as brotli 

9 except ImportError: 

10 import brotli 

11 

12 HAS_BROTLI = True 

13except ImportError: # pragma: no cover 

14 HAS_BROTLI = False 

15 

16MAX_SYNC_CHUNK_SIZE = 1024 

17 

18 

19def encoding_to_mode( 

20 encoding: Optional[str] = None, 

21 suppress_deflate_header: bool = False, 

22) -> int: 

23 if encoding == "gzip": 

24 return 16 + zlib.MAX_WBITS 

25 

26 return -zlib.MAX_WBITS if suppress_deflate_header else zlib.MAX_WBITS 

27 

28 

29class ZlibBaseHandler: 

30 def __init__( 

31 self, 

32 mode: int, 

33 executor: Optional[Executor] = None, 

34 max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, 

35 ): 

36 self._mode = mode 

37 self._executor = executor 

38 self._max_sync_chunk_size = max_sync_chunk_size 

39 

40 

41class ZLibCompressor(ZlibBaseHandler): 

42 def __init__( 

43 self, 

44 encoding: Optional[str] = None, 

45 suppress_deflate_header: bool = False, 

46 level: Optional[int] = None, 

47 wbits: Optional[int] = None, 

48 strategy: int = zlib.Z_DEFAULT_STRATEGY, 

49 executor: Optional[Executor] = None, 

50 max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, 

51 ): 

52 super().__init__( 

53 mode=encoding_to_mode(encoding, suppress_deflate_header) 

54 if wbits is None 

55 else wbits, 

56 executor=executor, 

57 max_sync_chunk_size=max_sync_chunk_size, 

58 ) 

59 if level is None: 

60 self._compressor = zlib.compressobj(wbits=self._mode, strategy=strategy) 

61 else: 

62 self._compressor = zlib.compressobj( 

63 wbits=self._mode, strategy=strategy, level=level 

64 ) 

65 

66 def compress_sync(self, data: bytes) -> bytes: 

67 return self._compressor.compress(data) 

68 

69 async def compress(self, data: bytes) -> bytes: 

70 if ( 

71 self._max_sync_chunk_size is not None 

72 and len(data) > self._max_sync_chunk_size 

73 ): 

74 return await asyncio.get_event_loop().run_in_executor( 

75 self._executor, self.compress_sync, data 

76 ) 

77 return self.compress_sync(data) 

78 

79 def flush(self, mode: int = zlib.Z_FINISH) -> bytes: 

80 return self._compressor.flush(mode) 

81 

82 

83class ZLibDecompressor(ZlibBaseHandler): 

84 def __init__( 

85 self, 

86 encoding: Optional[str] = None, 

87 suppress_deflate_header: bool = False, 

88 executor: Optional[Executor] = None, 

89 max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE, 

90 ): 

91 super().__init__( 

92 mode=encoding_to_mode(encoding, suppress_deflate_header), 

93 executor=executor, 

94 max_sync_chunk_size=max_sync_chunk_size, 

95 ) 

96 self._decompressor = zlib.decompressobj(wbits=self._mode) 

97 

98 def decompress_sync(self, data: bytes, max_length: int = 0) -> bytes: 

99 return self._decompressor.decompress(data, max_length) 

100 

101 async def decompress(self, data: bytes, max_length: int = 0) -> bytes: 

102 if ( 

103 self._max_sync_chunk_size is not None 

104 and len(data) > self._max_sync_chunk_size 

105 ): 

106 return await asyncio.get_event_loop().run_in_executor( 

107 self._executor, self.decompress_sync, data, max_length 

108 ) 

109 return self.decompress_sync(data, max_length) 

110 

111 def flush(self, length: int = 0) -> bytes: 

112 return ( 

113 self._decompressor.flush(length) 

114 if length > 0 

115 else self._decompressor.flush() 

116 ) 

117 

118 @property 

119 def eof(self) -> bool: 

120 return self._decompressor.eof 

121 

122 @property 

123 def unconsumed_tail(self) -> bytes: 

124 return self._decompressor.unconsumed_tail 

125 

126 @property 

127 def unused_data(self) -> bytes: 

128 return self._decompressor.unused_data 

129 

130 

131class BrotliDecompressor: 

132 # Supports both 'brotlipy' and 'Brotli' packages 

133 # since they share an import name. The top branches 

134 # are for 'brotlipy' and bottom branches for 'Brotli' 

135 def __init__(self) -> None: 

136 if not HAS_BROTLI: 

137 raise RuntimeError( 

138 "The brotli decompression is not available. " 

139 "Please install `Brotli` module" 

140 ) 

141 self._obj = brotli.Decompressor() 

142 

143 def decompress_sync(self, data: bytes) -> bytes: 

144 if hasattr(self._obj, "decompress"): 

145 return cast(bytes, self._obj.decompress(data)) 

146 return cast(bytes, self._obj.process(data)) 

147 

148 def flush(self) -> bytes: 

149 if hasattr(self._obj, "flush"): 

150 return cast(bytes, self._obj.flush()) 

151 return b""