import importlib

from .. import errors
from ..protocol import CompressionMethodByte


def get_compressor_cls(alg):
    try:
        module = importlib.import_module("." + alg, __name__)
        return module.Compressor

    except ImportError:
        raise errors.UnknownCompressionMethod("Unknown compression method: '{}'".format(alg))


def get_decompressor_cls(method_type):
    if method_type == CompressionMethodByte.LZ4:
        module = importlib.import_module(".lz4", __name__)

    elif method_type == CompressionMethodByte.ZSTD:
        module = importlib.import_module(".zstd", __name__)

    else:
        raise errors.UnknownCompressionMethod()

    return module.Decompressor
