tweaks.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

__author__ = "Jakub Barylski"
__maintainer__ = "Jakub Barylski"
__license__ = "GNU GENERAL PUBLIC LICENSE"
__email__ = "jakub.barylski@gmail.com"

import inspect
import pickle
import sys
import warnings
from functools import wraps
from multiprocessing import cpu_count
from pathlib import Path
from subprocess import run, DEVNULL
from typing import Callable, Dict, Collection, List, Hashable, Any

import joblib
import numpy as np
from loguru import logger
from tqdm import tqdm


# LOGGING format and configuration
def loguru_showwarning(message, category, filename, lineno, file=None, line=None):
    filename = Path(filename).name
    logger.warning(f"{filename} {category.__name__}: {message}")


logger.remove()
log_format = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> <cyan>{function}</cyan>: <level>{message}</level>"
logger.add(sys.stderr, format=log_format)
warnings.showwarning = loguru_showwarning

# DEFAULTS
default_threads = max(cpu_count() - 1, 1)
seed_max = 2 ** 32 - 1

extensions = {'gbk': {'.gb', '.gbk'},
              'gff': {'.gff', '.gff3'},
              'fna': {'.fa', '.fas', '.fasta', '.fna'},
              'faa': {'.fa', '.fas', '.fasta', '.faa'}}

for form in extensions:  # add gzip compressed files
    extensions[form].update([f'{e}.gz' for e in extensions[form]])


def frantic_search(dictionary: Dict[Hashable, Any],
                   *possible_keys: Hashable):
    """
    Find first of the several keys that are in a dictionary and return the value
    :param dictionary: dictionary to search
    :param possible_keys: any number of possible keys (preferred first)
    :return: dictionary[first_found_key]
    """
    for key in possible_keys:
        if key in dictionary:
            return dictionary[key]
    missed_keys = ', '.join([str(k) for k in possible_keys])
    raise KeyError(f'Found none of the: {missed_keys}')


def find_files(directory: Path,
               file_type: str,
               descent: bool = False) -> List[Path]:
    """
    Find files of a given type in a folder.
    :param directory: directory to search in
    :param file_type: file type to search for
    :param descent: whether to search in subdirectories (and their subdirectories)
    :return:
    """
    main_file_catalogue = []
    expected_extensions = extensions[file_type]
    detected_files = [f for f in directory.iterdir() if f.suffix in expected_extensions]
    subdirectories = [f for f in directory.iterdir() if f.is_dir()]

    if detected_files:
        main_file_catalogue.extend(detected_files)

    elif subdirectories and descent:
        for sd in subdirectories:
            main_file_catalogue.extend(find_files(sd, file_type, descent))

    return main_file_catalogue


# PARALLELIZATION
class Parallel(joblib.Parallel):
    """
    The modification of joblib.Parallel
    with a TQDM progress bar
    according to Nth
    (https://stackoverflow.com/questions/37804279/how-can-we-use-tqdm-in-a-parallel-execution-with-joblib)
    """

    def __init__(self,
                 parallelized_function: Callable,
                 input_collection: Collection = None,
                 random_replicates: int = None,
                 kwargs: Dict = None,
                 n_jobs=None,
                 backend=None,
                 description: str = None,
                 verbose=0,
                 timeout=None,
                 pre_dispatch='2 * n_jobs',
                 batch_size='auto',
                 temp_folder=None, max_nbytes='1M', mmap_mode='r',
                 prefer=None,
                 require=None,
                 bar: bool = True):

        if not n_jobs:
            n_jobs = default_threads

        joblib.Parallel.__init__(self,
                                 n_jobs=n_jobs,
                                 backend=backend,
                                 verbose=verbose,
                                 timeout=timeout,
                                 pre_dispatch=pre_dispatch,
                                 batch_size=batch_size,
                                 temp_folder=temp_folder,
                                 max_nbytes=max_nbytes,
                                 mmap_mode=mmap_mode,
                                 prefer=prefer,
                                 require=require)

        assert bool(random_replicates) ^ bool(input_collection), 'You need to specify EITHER an input collection ' \
                                                                 'OR number of random replicates of the function'

        kwargs = {} if not kwargs else kwargs

        description = description if description else parallelized_function.__name__

        if random_replicates:
            input_collection = np.random.randint(0, 2 ** 32 - 1,
                                                 size=random_replicates,
                                                 dtype=np.int64)
            description = f'{description} 🎲'

        jobs = ((joblib.delayed(parallelized_function)(e, **kwargs)) for e in input_collection)

        if bar:
            self._progress = tqdm(total=len(input_collection), file=sys.stdout)
            if description:
                self._progress.set_description(description)
        else:
            self._progress = None

        self.result = list(self.__call__(jobs))

        if self._progress:
            self._progress.close()
            print(flush=True)

    def print_progress(self):
        if self._progress:
            self._progress.n = self.n_completed_tasks
            self._progress.refresh()


def run_external(command: List[str],
                 stdout='suppress',
                 stdin=None):
    """
    Run external (non-python) command
    :param command: list of the expressions
                    that make up the shell command e.g. ['ls', '-lh']
    :param stdout: do not print the log messages
    :param stdin: input for the command
    """
    sanitized_command = [str(c) for c in command]

    logger.info(" ".join(sanitized_command))
    if stdout == 'suppress':
        process = run(sanitized_command, stdout=DEVNULL, stderr=DEVNULL, input=stdin)
    elif stdout == 'capture':
        process = run(sanitized_command, capture_output=True, input=stdin)
        return process.stdout
    else:
        process = run(sanitized_command)
    if process.returncode or process.stderr:
        raise ChildProcessError(f'"{" ".join(sanitized_command)}" crashed with:\n{process.stderr}')


def parse_fasta(fasta: Path):
    """
    Simple and relatively fast fasta parser
    used when no complex sequence handling is required
    :param fasta: path to a fasta file
    :return: generator yielding tuples of (identifier, sequence)
    """
    identifier, sequence = None, []
    with fasta.open() as fas:
        for line in fas:
            line = line.rstrip('\n')
            if line.startswith('>'):
                if identifier is not None:
                    yield identifier, ''.join(sequence)
                identifier = line.lstrip('>').split(' ')[0]
                sequence = []
            else:
                sequence.append(line)
    yield identifier, ''.join(sequence)


def checkpoint(funct: callable):
    """
    Simple serialization decorator
    that saves the function result
    if exacted output file doesn't exist or is empty
    or read it if it is non-empty
    @param funct: function to wrap
    @param pickle_path: a path to an output file
    @param serialization_method: a module used for serialization (either joblib or pickle)
    @return:
    """

    signature = inspect.signature(funct)

    @wraps(funct)
    def save_checkpoint(*args, **kwargs):

        bound_args = signature.bind(*args, **kwargs)
        pickle_path = Path(bound_args.arguments.get('pickle_path',
                                                    signature.parameters['pickle_path'].default))
        if pickle_path:
            try:
                with pickle_path.open('rb') as file_object:
                    result = pickle.load(file_object)
                logger.info(f'\ntemporary file read from: {pickle_path.as_posix()}\n', flush=True)
                return result
            except (FileNotFoundError, IOError, EOFError):
                sys.setrecursionlimit(5000)
                result = funct(*args, **kwargs)
                with pickle_path.open('wb') as out:
                    pickle.dump(result, out)
                logger.info(f'\ntemporary file stored at: {pickle_path.as_posix()}\n', flush=True)
                return result

    return save_checkpoint