Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/zstandard/__init__.py: 25%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

83 statements  

1# Copyright (c) 2017-present, Gregory Szorc 

2# All rights reserved. 

3# 

4# This software may be modified and distributed under the terms 

5# of the BSD license. See the LICENSE file for details. 

6 

7# ruff: noqa: F403, F405 

8 

9"""Python interface to the Zstandard (zstd) compression library.""" 

10 

11from __future__ import absolute_import, unicode_literals 

12 

13# This module serves 2 roles: 

14# 

15# 1) Export the C or CFFI "backend" through a central module. 

16# 2) Implement additional functionality built on top of C or CFFI backend. 

17import builtins 

18import io 

19import os 

20import platform 

21import sys 

22 

23if sys.version_info >= (3, 12): 

24 from collections.abc import Buffer 

25else: 

26 from typing import ByteString as Buffer 

27 

28# Some Python implementations don't support C extensions. That's why we have 

29# a CFFI implementation in the first place. The code here import one of our 

30# "backends" then re-exports the symbols from this module. For convenience, 

31# we support falling back to the CFFI backend if the C extension can't be 

32# imported. But for performance reasons, we only do this on unknown Python 

33# implementation. Notably, for CPython we require the C extension by default. 

34# Because someone will inevitably want special behavior, the behavior is 

35# configurable via an environment variable. A potentially better way to handle 

36# this is to import a special ``__importpolicy__`` module or something 

37# defining a variable and `setup.py` could write the file with whatever 

38# policy was specified at build time. Until someone needs it, we go with 

39# the hacky but simple environment variable approach. 

40_module_policy = os.environ.get( 

41 "PYTHON_ZSTANDARD_IMPORT_POLICY", "default" 

42).strip() 

43 

44if _module_policy == "default": 

45 if platform.python_implementation() in ("CPython",): 

46 from .backend_c import * # type: ignore 

47 

48 backend = "cext" 

49 elif platform.python_implementation() in ("PyPy",): 

50 from .backend_cffi import * # type: ignore 

51 

52 backend = "cffi" 

53 else: 

54 try: 

55 from .backend_c import * 

56 

57 backend = "cext" 

58 except ImportError: 

59 from .backend_cffi import * 

60 

61 backend = "cffi" 

62elif _module_policy == "cffi_fallback": 

63 try: 

64 from .backend_c import * 

65 

66 backend = "cext" 

67 except ImportError: 

68 from .backend_cffi import * 

69 

70 backend = "cffi" 

71elif _module_policy == "rust": 

72 from .backend_rust import * # type: ignore 

73 

74 backend = "rust" 

75elif _module_policy == "cext": 

76 from .backend_c import * 

77 

78 backend = "cext" 

79elif _module_policy == "cffi": 

80 from .backend_cffi import * 

81 

82 backend = "cffi" 

83else: 

84 raise ImportError( 

85 "unknown module import policy: %s; use default, cffi_fallback, " 

86 "cext, or cffi" % _module_policy 

87 ) 

88 

89# Keep this in sync with python-zstandard.h, rust-ext/src/lib.rs, and debian/changelog. 

90__version__ = "0.25.0" 

91 

92_MODE_CLOSED = 0 

93_MODE_READ = 1 

94_MODE_WRITE = 2 

95 

96 

97def open( 

98 filename, 

99 mode="rb", 

100 cctx=None, 

101 dctx=None, 

102 encoding=None, 

103 errors=None, 

104 newline=None, 

105 closefd=None, 

106): 

107 """Create a file object with zstd (de)compression. 

108 

109 The object returned from this function will be a 

110 :py:class:`ZstdDecompressionReader` if opened for reading in binary mode, 

111 a :py:class:`ZstdCompressionWriter` if opened for writing in binary mode, 

112 or an ``io.TextIOWrapper`` if opened for reading or writing in text mode. 

113 

114 :param filename: 

115 ``bytes``, ``str``, or ``os.PathLike`` defining a file to open or a 

116 file object (with a ``read()`` or ``write()`` method). 

117 :param mode: 

118 ``str`` File open mode. Accepts any of the open modes recognized by 

119 ``open()``. 

120 :param cctx: 

121 ``ZstdCompressor`` to use for compression. If not specified and file 

122 is opened for writing, the default ``ZstdCompressor`` will be used. 

123 :param dctx: 

124 ``ZstdDecompressor`` to use for decompression. If not specified and file 

125 is opened for reading, the default ``ZstdDecompressor`` will be used. 

126 :param encoding: 

127 ``str`` that defines text encoding to use when file is opened in text 

128 mode. 

129 :param errors: 

130 ``str`` defining text encoding error handling mode. 

131 :param newline: 

132 ``str`` defining newline to use in text mode. 

133 :param closefd: 

134 ``bool`` whether to close the file when the returned object is closed. 

135 Only used if a file object is passed. If a filename is specified, the 

136 opened file is always closed when the returned object is closed. 

137 """ 

138 normalized_mode = mode.replace("t", "") 

139 

140 if normalized_mode in ("r", "rb"): 

141 dctx = dctx or ZstdDecompressor() 

142 open_mode = "r" 

143 raw_open_mode = "rb" 

144 elif normalized_mode in ("w", "wb", "a", "ab", "x", "xb"): 

145 cctx = cctx or ZstdCompressor() 

146 open_mode = "w" 

147 raw_open_mode = normalized_mode 

148 if not raw_open_mode.endswith("b"): 

149 raw_open_mode = raw_open_mode + "b" 

150 else: 

151 raise ValueError("Invalid mode: {!r}".format(mode)) 

152 

153 if hasattr(os, "PathLike"): 

154 types = (str, bytes, os.PathLike) 

155 else: 

156 types = (str, bytes) 

157 

158 if isinstance(filename, types): # type: ignore 

159 inner_fh = builtins.open(filename, raw_open_mode) 

160 closefd = True 

161 elif hasattr(filename, "read") or hasattr(filename, "write"): 

162 inner_fh = filename 

163 closefd = bool(closefd) 

164 else: 

165 raise TypeError( 

166 "filename must be a str, bytes, file or PathLike object" 

167 ) 

168 

169 if open_mode == "r": 

170 fh = dctx.stream_reader(inner_fh, closefd=closefd) 

171 elif open_mode == "w": 

172 fh = cctx.stream_writer(inner_fh, closefd=closefd) 

173 else: 

174 raise RuntimeError("logic error in zstandard.open() handling open mode") 

175 

176 if "b" not in normalized_mode: 

177 return io.TextIOWrapper( 

178 fh, encoding=encoding, errors=errors, newline=newline 

179 ) 

180 else: 

181 return fh 

182 

183 

184def compress(data: Buffer, level: int = 3) -> bytes: 

185 """Compress source data using the zstd compression format. 

186 

187 This performs one-shot compression using basic/default compression 

188 settings. 

189 

190 This method is provided for convenience and is equivalent to calling 

191 ``ZstdCompressor(level=level).compress(data)``. 

192 

193 If you find yourself calling this function in a tight loop, 

194 performance will be greater if you construct a single ``ZstdCompressor`` 

195 and repeatedly call ``compress()`` on it. 

196 """ 

197 cctx = ZstdCompressor(level=level) 

198 

199 return cctx.compress(data) 

200 

201 

202def decompress(data: Buffer, max_output_size: int = 0) -> bytes: 

203 """Decompress a zstd frame into its original data. 

204 

205 This performs one-shot decompression using basic/default compression 

206 settings. 

207 

208 This method is provided for convenience and is equivalent to calling 

209 ``ZstdDecompressor().decompress(data, max_output_size=max_output_size)``. 

210 

211 If you find yourself calling this function in a tight loop, performance 

212 will be greater if you construct a single ``ZstdDecompressor`` and 

213 repeatedly call ``decompress()`` on it. 

214 """ 

215 dctx = ZstdDecompressor() 

216 

217 return dctx.decompress(data, max_output_size=max_output_size)