def test_register_compressor_invalid_fileobj():
# Test that registering an invalid file object is not allowed.
class InvalidFileObject():
pass
class InvalidFileObjectWrapper(CompressorWrapper):
def __init__(self):
CompressorWrapper.__init__(self, obj=InvalidFileObject,
prefix=b'prefix')
with raises(ValueError) as excinfo:
register_compressor('invalid', InvalidFileObjectWrapper())
excinfo.match("Compressor 'fileobj_factory' attribute should implement "
"the file object interface")
def test_register_compressor(tmpdir):
# Check that registering compressor file works.
compressor_name = 'test-name'
compressor_prefix = 'test-prefix'
class BinaryCompressorTestFile(io.BufferedIOBase):
pass
class BinaryCompressorTestWrapper(CompressorWrapper):
def __init__(self):
CompressorWrapper.__init__(self,
obj=BinaryCompressorTestFile,
prefix=compressor_prefix)
register_compressor(compressor_name, BinaryCompressorTestWrapper())
assert (_COMPRESSORS[compressor_name].fileobj_factory ==
BinaryCompressorTestFile)
assert _COMPRESSORS[compressor_name].prefix == compressor_prefix
# Remove this dummy compressor file from extra compressors because other
# tests might fail because of this
_COMPRESSORS.pop(compressor_name)
def test_register_compressor_already_registered():
# Test registration of existing compressor files.
compressor_name = 'test-name'
# register a test compressor
register_compressor(compressor_name, AnotherZlibCompressorWrapper())
with raises(ValueError) as excinfo:
register_compressor(compressor_name,
StandardLibGzipCompressorWrapper())
excinfo.match("Compressor '{}' already registered."
.format(compressor_name))
register_compressor(compressor_name, StandardLibGzipCompressorWrapper(),
force=True)
assert compressor_name in _COMPRESSORS
assert _COMPRESSORS[compressor_name].fileobj_factory == gzip.GzipFile
# Remove this dummy compressor file from extra compressors because other
# tests might fail because of this
_COMPRESSORS.pop(compressor_name)
def test_register_compressor_invalid_name(invalid_name):
# Test that registering an invalid compressor name is not allowed.
with raises(ValueError) as excinfo:
register_compressor(invalid_name, None)
excinfo.match("Compressor name should be a string")
def register(compress_levels=None, compressor_args=None, decompressor_args=None, force=True):
if compress_levels and (not isinstance(compress_levels, (list, tuple)) or len(compress_levels) != 10):
raise ValueError("compress_levels L must be list or tuple, len(L) == 10, 1 <= L[i] <= 22.")
joblib.register_compressor(
"zstd", ZStandardCompressorWrapper(compress_levels, compressor_args, decompressor_args), force=force,
)
