1616import warnings
1717from typing import Any , Iterable , Optional , Union
1818
19- from pymongo ._lazy_import import lazy_import
2019from pymongo .hello import HelloCompat
21- from pymongo .monitoring import _SENSITIVE_COMMANDS
20+ from pymongo .helpers import _SENSITIVE_COMMANDS
2221
23- try :
24- snappy = lazy_import ("snappy" )
25- _HAVE_SNAPPY = True
26- except ImportError :
27- # python-snappy isn't available.
28- _HAVE_SNAPPY = False
22+ _SUPPORTED_COMPRESSORS = {"snappy" , "zlib" , "zstd" }
23+ _NO_COMPRESSION = {HelloCompat .CMD , HelloCompat .LEGACY_CMD }
24+ _NO_COMPRESSION .update (_SENSITIVE_COMMANDS )
2925
30- try :
31- zlib = lazy_import ("zlib" )
3226
33- _HAVE_ZLIB = True
34- except ImportError :
35- # Python built without zlib support.
36- _HAVE_ZLIB = False
27+ def _have_snappy () -> bool :
28+ try :
29+ import snappy # type:ignore[import] # noqa: F401
3730
38- try :
39- zstandard = lazy_import ("zstandard" )
40- _HAVE_ZSTD = True
41- except ImportError :
42- _HAVE_ZSTD = False
31+ return True
32+ except ImportError :
33+ return False
4334
44- _SUPPORTED_COMPRESSORS = {"snappy" , "zlib" , "zstd" }
45- _NO_COMPRESSION = {HelloCompat .CMD , HelloCompat .LEGACY_CMD }
46- _NO_COMPRESSION .update (_SENSITIVE_COMMANDS )
35+
36+ def _have_zlib () -> bool :
37+ try :
38+ import zlib # noqa: F401
39+
40+ return True
41+ except ImportError :
42+ return False
43+
44+
45+ def _have_zstd () -> bool :
46+ try :
47+ import zstandard # noqa: F401
48+
49+ return True
50+ except ImportError :
51+ return False
4752
4853
4954def validate_compressors (dummy : Any , value : Union [str , Iterable [str ]]) -> list [str ]:
@@ -58,21 +63,21 @@ def validate_compressors(dummy: Any, value: Union[str, Iterable[str]]) -> list[s
5863 if compressor not in _SUPPORTED_COMPRESSORS :
5964 compressors .remove (compressor )
6065 warnings .warn (f"Unsupported compressor: { compressor } " , stacklevel = 2 )
61- elif compressor == "snappy" and not _HAVE_SNAPPY :
66+ elif compressor == "snappy" and not _have_snappy () :
6267 compressors .remove (compressor )
6368 warnings .warn (
6469 "Wire protocol compression with snappy is not available. "
6570 "You must install the python-snappy module for snappy support." ,
6671 stacklevel = 2 ,
6772 )
68- elif compressor == "zlib" and not _HAVE_ZLIB :
73+ elif compressor == "zlib" and not _have_zlib () :
6974 compressors .remove (compressor )
7075 warnings .warn (
7176 "Wire protocol compression with zlib is not available. "
7277 "The zlib module is not available." ,
7378 stacklevel = 2 ,
7479 )
75- elif compressor == "zstd" and not _HAVE_ZSTD :
80+ elif compressor == "zstd" and not _have_zstd () :
7681 compressors .remove (compressor )
7782 warnings .warn (
7883 "Wire protocol compression with zstandard is not available. "
@@ -117,6 +122,8 @@ class SnappyContext:
117122
118123 @staticmethod
119124 def compress (data : bytes ) -> bytes :
125+ import snappy
126+
120127 return snappy .compress (data )
121128
122129
@@ -127,6 +134,8 @@ def __init__(self, level: int):
127134 self .level = level
128135
129136 def compress (self , data : bytes ) -> bytes :
137+ import zlib
138+
130139 return zlib .compress (data , self .level )
131140
132141
@@ -137,6 +146,8 @@ class ZstdContext:
137146 def compress (data : bytes ) -> bytes :
138147 # ZstdCompressor is not thread safe.
139148 # TODO: Use a pool?
149+ import zstandard
150+
140151 return zstandard .ZstdCompressor ().compress (data )
141152
142153
@@ -146,12 +157,18 @@ def decompress(data: bytes, compressor_id: int) -> bytes:
146157 # https://github.com/andrix/python-snappy/issues/65
147158 # This only matters when data is a memoryview since
148159 # id(bytes(data)) == id(data) when data is a bytes.
160+ import snappy
161+
149162 return snappy .uncompress (bytes (data ))
150163 elif compressor_id == ZlibContext .compressor_id :
164+ import zlib
165+
151166 return zlib .decompress (data )
152167 elif compressor_id == ZstdContext .compressor_id :
153168 # ZstdDecompressor is not thread safe.
154169 # TODO: Use a pool?
170+ import zstandard
171+
155172 return zstandard .ZstdDecompressor ().decompress (data )
156173 else :
157174 raise ValueError ("Unknown compressorId %d" % (compressor_id ,))
0 commit comments