Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/smart_open/compression.py: 40%

1# -*- coding: utf-8 -*- 

2# 

3# Copyright (C) 2020 Radim Rehurek <me@radimrehurek.com> 

4# 

5# This code is distributed under the terms and conditions 

6# from the MIT License (MIT). 

7# 

8"""Implements the compression layer of the ``smart_open`` library.""" 

9import io 

10import logging 

11import os.path 

12 

13logger = logging.getLogger(__name__) 

14 

15_COMPRESSOR_REGISTRY = {} 

16 

17NO_COMPRESSION = 'disable' 

18"""Use no compression. Read/write the data as-is.""" 

19INFER_FROM_EXTENSION = 'infer_from_extension' 

20"""Determine the compression to use from the file extension. 

21 

22See get_supported_extensions(). 

23""" 

24 

25 

26def get_supported_compression_types(): 

27 """Return the list of supported compression types available to open. 

28 

29 See compression paratemeter to smart_open.open(). 

30 """ 

31 return [NO_COMPRESSION, INFER_FROM_EXTENSION] + get_supported_extensions() 

32 

33 

34def get_supported_extensions(): 

35 """Return the list of file extensions for which we have registered compressors.""" 

36 return sorted(_COMPRESSOR_REGISTRY.keys()) 

37 

38 

39def register_compressor(ext, callback): 

40 """Register a callback for transparently decompressing files with a specific extension. 

41 

42 Parameters 

43 ---------- 

44 ext: str 

45 The extension. Must include the leading period, e.g. ``.gz``. 

46 callback: callable 

47 The callback. It must accept two position arguments, file_obj and mode. 

48 This function will be called when ``smart_open`` is opening a file with 

49 the specified extension. 

50 

51 Examples 

52 -------- 

53 

54 Instruct smart_open to use the `lzma` module whenever opening a file 

55 with a .xz extension (see README.rst for the complete example showing I/O): 

56 

57 >>> def _handle_xz(file_obj, mode): 

58 ... import lzma 

59 ... return lzma.LZMAFile(filename=file_obj, mode=mode, format=lzma.FORMAT_XZ) 

60 >>> 

61 >>> register_compressor('.xz', _handle_xz) 

62 

63 """ 

64 if not (ext and ext[0] == '.'): 

65 raise ValueError('ext must be a string starting with ., not %r' % ext) 

66 ext = ext.lower() 

67 if ext in _COMPRESSOR_REGISTRY: 

68 logger.warning('overriding existing compression handler for %r', ext) 

69 _COMPRESSOR_REGISTRY[ext] = callback 

70 

71 

72def tweak_close(outer, inner): 

73 """Ensure that closing the `outer` stream closes the `inner` stream as well. 

74 

75 Deprecated: smart_open.open().__exit__ now always calls __exit__ on the 

76 underlying filestream. 

77 

78 Use this when your compression library's `close` method does not 

79 automatically close the underlying filestream. See 

80 https://github.com/piskvorky/smart_open/issues/630 for an 

81 explanation why that is a problem for smart_open. 

82 """ 

83 outer_close = outer.close 

84 

85 def close_both(*args): 

86 nonlocal inner 

87 try: 

88 outer_close() 

89 finally: 

90 if inner: 

91 inner, fp = None, inner 

92 fp.close() 

93 

94 outer.close = close_both 

95 

96 

97def _handle_bz2(file_obj, mode): 

98 from bz2 import BZ2File 

99 result = BZ2File(file_obj, mode) 

100 return result 

101 

102 

103def _handle_gzip(file_obj, mode): 

104 import gzip 

105 result = gzip.GzipFile(fileobj=file_obj, mode=mode) 

106 return result 

107 

108 

109def _handle_zstd(file_obj, mode): 

110 import zstandard # type: ignore 

111 result = zstandard.open(filename=file_obj, mode=mode) 

112 # zstandard.open returns an io.TextIOWrapper in text mode, but otherwise 

113 # returns a raw stream reader/writer, and we need the `io` wrapper 

114 # to make FileLikeProxy work correctly. 

115 # 

116 # See: 

117 # 

118 # https://github.com/indygreg/python-zstandard/blob/d7d81e79dbe74feb22fb73405ebfb3e20f4c4653/zstandard/__init__.py#L169-L174 

119 if "b" in mode and "w" in mode: 

120 result = io.BufferedWriter(result) 

121 elif "b" in mode and "r" in mode: 

122 result = io.BufferedReader(result) 

123 return result 

124 

125 

126def compression_wrapper(file_obj, mode, compression=INFER_FROM_EXTENSION, filename=None): 

127 """ 

128 Wrap `file_obj` with an appropriate [de]compression mechanism based on its file extension. 

129 

130 If the filename extension isn't recognized, simply return the original `file_obj` unchanged. 

131 

132 `file_obj` must either be a filehandle object, or a class which behaves like one. 

133 

134 If `filename` is specified, it will be used to extract the extension. 

135 If not, the `file_obj.name` attribute is used as the filename. 

136 

137 """ 

138 if compression == NO_COMPRESSION: 

139 return file_obj 

140 elif compression == INFER_FROM_EXTENSION: 

141 try: 

142 filename = (filename or file_obj.name).lower() 

143 except (AttributeError, TypeError): 

144 logger.warning( 

145 'unable to transparently decompress %r because it ' 

146 'seems to lack a string-like .name', file_obj 

147 ) 

148 return file_obj 

149 _, compression = os.path.splitext(filename) 

150 

151 if compression in _COMPRESSOR_REGISTRY and mode.endswith('+'): 

152 raise ValueError('transparent (de)compression unsupported for mode %r' % mode) 

153 

154 try: 

155 callback = _COMPRESSOR_REGISTRY[compression] 

156 except KeyError: 

157 return file_obj 

158 else: 

159 return callback(file_obj, mode) 

160 

161 

162# 

163# NB. avoid using lambda here to make stack traces more readable. 

164# 

165register_compressor('.bz2', _handle_bz2) 

166register_compressor('.gz', _handle_gzip) 

167register_compressor('.zst', _handle_zstd)