codekingpro/portable-devtools

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)
218