diff --git a/fimfarchive/fetchers/fimfarchive.py b/fimfarchive/fetchers/fimfarchive.py index 8da3841..b704696 100644 --- a/fimfarchive/fetchers/fimfarchive.py +++ b/fimfarchive/fetchers/fimfarchive.py @@ -32,17 +32,11 @@ from jmespath import compile as jmes from fimfarchive.exceptions import InvalidStoryError, StorySourceError from fimfarchive.flavors import StorySource, DataFormat, MetaPurity from fimfarchive.stories import Story +from fimfarchive.utils import find_compressor from .base import Fetcher -try: - from lz4.block import compress, decompress -except ModuleNotFoundError: - compress = lambda data: data # noqa - decompress = lambda data: data # noqa - - __all__ = ( 'FimfarchiveFetcher', ) @@ -52,6 +46,9 @@ BUFFER_SIZE = 8_000_000 PATH = jmes('archive.path || path') +compress, decompress = find_compressor() + + class FimfarchiveFetcher(Fetcher): """ Fetcher for Fimfarchive. @@ -76,7 +73,7 @@ class FimfarchiveFetcher(Fetcher): StorySourceError: If no valid Fimfarchive release can be loaded. """ self.archive: ZipFile - self.index: Dict[int, str] + self.index: Dict[int, bytes] self.paths: Dict[int, str] self.is_open: bool = False @@ -115,7 +112,10 @@ class FimfarchiveFetcher(Fetcher): self.paths = dict() self.is_open = True - def load_index(self, source: Iterator[bytes]) -> Iterator[Tuple[int, str]]: + def load_index( + self, + source: Iterator[bytes], + ) -> Iterator[Tuple[int, bytes]]: """ Yields unparsed index items from a byte stream. diff --git a/fimfarchive/utils.py b/fimfarchive/utils.py index 763f8d7..7194d93 100644 --- a/fimfarchive/utils.py +++ b/fimfarchive/utils.py @@ -26,8 +26,12 @@ import json import os import shutil from functools import partial +from importlib import import_module from importlib_resources import read_binary, read_text -from typing import Any, Dict, Iterator, Optional, Type, TypeVar, Union +from typing import ( + cast, Any, Callable, Dict, Iterator, + Optional, Tuple, Type, TypeVar, Union, +) from tqdm import tqdm @@ -38,12 +42,14 @@ from fimfarchive.stories import Story __all__ = ( 'Empty', 'PersistedDict', + 'find_compressor', 'find_flavor', 'tqdm', ) F = TypeVar('F', bound=Flavor) +ByteFunc = Callable[[bytes], bytes] tqdm = partial( @@ -160,6 +166,25 @@ class JayWalker: self.walk(value) +def find_compressor() -> Tuple[ByteFunc, ByteFunc]: + """ + Searches for a fast compression module. + + Returns: + A pair of compression functions. + """ + for compressor in ('lz4.block', 'snappy', 'lzo'): + try: + module: Any = import_module(compressor) + return module.compress, module.decompress + except ImportError: + pass + + dummy = cast(ByteFunc, lambda data: data) + + return dummy, dummy + + def find_flavor(story: Story, flavor: Type[F]) -> Optional[F]: """ Searches for a flavor of a specific type.