# Copyright 2018 MongoDB, Inc.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from __future__ import annotations

import sys

import warnings

from collections.abc import Iterable

from typing import Any, Optional, Union

from pymongo.hello import HelloCompat

from pymongo.helpers_shared import _SENSITIVE_COMMANDS

_SUPPORTED_COMPRESSORS = {"snappy", "zlib", "zstd"}

_NO_COMPRESSION = {HelloCompat.CMD, HelloCompat.LEGACY_CMD}

_NO_COMPRESSION.update(_SENSITIVE_COMMANDS)

def _have_snappy() -> bool:

try:

import snappy # type:ignore[import-untyped] # noqa: F401

return True

except ImportError:

return False

def _have_zlib() -> bool:

try:

import zlib # noqa: F401

return True

except ImportError:

return False

def _have_zstd() -> bool:

try:

if sys.version_info >= (3, 14):

from compression import zstd

else:

from backports import zstd # noqa: F401

return True

except ImportError:

return False

def validate_compressors(dummy: Any, value: Union[str, Iterable[str]]) -> list[str]:

try:

# `value` is string.

compressors = value.split(",") # type: ignore[union-attr]

except AttributeError:

# `value` is an iterable.

compressors = list(value)

for compressor in compressors[:]:

if compressor not in _SUPPORTED_COMPRESSORS:

compressors.remove(compressor)

warnings.warn(f"Unsupported compressor: {compressor}", stacklevel=2)

elif compressor == "snappy" and not _have_snappy():

compressors.remove(compressor)

warnings.warn(

"Wire protocol compression with snappy is not available. "

"You must install the python-snappy module for snappy support.",

stacklevel=2,

)

elif compressor == "zlib" and not _have_zlib():

compressors.remove(compressor)

warnings.warn(

"Wire protocol compression with zlib is not available. "

"The zlib module is not available.",

stacklevel=2,

)

elif compressor == "zstd" and not _have_zstd():

compressors.remove(compressor)

if sys.version_info >= (3, 14):

warnings.warn(

"Wire protocol compression with zstandard is not available. "

"The compression.zstd module is not available.",

stacklevel=2,

)

else:

warnings.warn(

"Wire protocol compression with zstandard is not available. "

"You must install the backports.zstd module for zstandard support.",

stacklevel=2,

)

return compressors

def validate_zlib_compression_level(option: str, value: Any) -> int:

try:

level = int(value)

except Exception:

raise TypeError(f"{option} must be an integer, not {value!r}") from None

if level < -1 or level > 9:

raise ValueError(f"{option} must be between -1 and 9, not {level}.")

return level

class CompressionSettings:

def __init__(self, compressors: list[str], zlib_compression_level: int):

self.compressors = compressors

self.zlib_compression_level = zlib_compression_level

def get_compression_context(

self, compressors: Optional[list[str]]

) -> Union[SnappyContext, ZlibContext, ZstdContext, None]:

if compressors:

chosen = compressors[0]

if chosen == "snappy":

return SnappyContext()

elif chosen == "zlib":

return ZlibContext(self.zlib_compression_level)

elif chosen == "zstd":

return ZstdContext()

return None

return None

class SnappyContext:

compressor_id = 1

@staticmethod

def compress(data: bytes) -> bytes:

import snappy

return snappy.compress(data)

class ZlibContext:

compressor_id = 2

def __init__(self, level: int):

self.level = level

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

import zlib

return zlib.compress(data, self.level)

class ZstdContext:

compressor_id = 3

@staticmethod

def compress(data: bytes) -> bytes:

if sys.version_info >= (3, 14):

from compression import zstd

else:

from backports import zstd

return zstd.compress(data)

def decompress(data: bytes | memoryview, compressor_id: int) -> bytes:

if compressor_id == SnappyContext.compressor_id:

# python-snappy doesn't support the buffer interface.

# https://github.com/andrix/python-snappy/issues/65

# This only matters when data is a memoryview since

# id(bytes(data)) == id(data) when data is a bytes.

import snappy

return snappy.uncompress(bytes(data))

elif compressor_id == ZlibContext.compressor_id:

import zlib

return zlib.decompress(data)

elif compressor_id == ZstdContext.compressor_id:

if sys.version_info >= (3, 14):

from compression import zstd

else:

from backports import zstd

return zstd.decompress(data)

else:

raise ValueError(f"Unknown compressorId {compressor_id}")