dupeguru/hscommon/util.py

# Created By: Virgil Dupras
# Created On: 2011-01-11
# Copyright 2015 Hardcoded Software (http://www.hardcoded.net)
#
# This software is licensed under the "GPLv3" License as described in the "LICENSE" file,
# which should be included with this package. The terms are also available at
# http://www.gnu.org/licenses/gpl-3.0.html

from math import ceil
from pathlib import Path
from hscommon.path import pathify, log_io_error

from typing import IO, Any, Callable, Generator, Iterable, List, Tuple, Union


def nonone(value: Any, replace_value: Any) -> Any:
    """Returns ``value`` if ``value`` is not ``None``. Returns ``replace_value`` otherwise."""
    if value is None:
        return replace_value
    else:
        return value


def tryint(value: Any, default: int = 0) -> int:
    """Tries to convert ``value`` to in ``int`` and returns ``default`` if it fails."""
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


# --- Sequence related


def dedupe(iterable: Iterable[Any]) -> List[Any]:
    """Returns a list of elements in ``iterable`` with all dupes removed.

    The order of the elements is preserved.
    """
    result = []
    seen = {}
    for item in iterable:
        if item in seen:
            continue
        seen[item] = 1
        result.append(item)
    return result


def flatten(iterables: Iterable[Iterable], start_with: Iterable[Any] = None) -> List[Any]:
    """Takes a list of lists ``iterables`` and returns a list containing elements of every list.

    If ``start_with`` is not ``None``, the result will start with ``start_with`` items, exactly as
    if ``start_with`` would be the first item of lists.
    """
    result: List[Any] = []
    if start_with:
        result.extend(start_with)
    for iterable in iterables:
        result.extend(iterable)
    return result


def first(iterable: Iterable[Any]):
    """Returns the first item of ``iterable``."""
    try:
        return next(iter(iterable))
    except StopIteration:
        return None


def extract(predicate: Callable[[Any], bool], iterable: Iterable[Any]) -> Tuple[List[Any], List[Any]]:
    """Separates the wheat from the shaft (`predicate` defines what's the wheat), and returns both."""
    wheat = []
    shaft = []
    for item in iterable:
        if predicate(item):
            wheat.append(item)
        else:
            shaft.append(item)
    return wheat, shaft


def allsame(iterable: Iterable[Any]) -> bool:
    """Returns whether all elements of 'iterable' are the same."""
    it = iter(iterable)
    try:
        first_item = next(it)
    except StopIteration:
        raise ValueError("iterable cannot be empty")
    return all(element == first_item for element in it)


def iterconsume(seq: List[Any], reverse: bool = True) -> Generator[Any, None, None]:
    """Iterate over ``seq`` and pops yielded objects.

    Because we use the ``pop()`` method, we reverse ``seq`` before proceeding. If you don't need
    to do that, set ``reverse`` to ``False``.

    This is useful in tight memory situation where you are looping over a sequence of objects that
    are going to be discarded afterwards. If you're creating other objects during that iteration
    you might want to use this to avoid ``MemoryError``.
    """
    if reverse:
        seq.reverse()
    while seq:
        yield seq.pop()


# --- String related


def escape(s: str, to_escape: str, escape_with: str = "\\") -> str:
    """Returns ``s`` with characters in ``to_escape`` all prepended with ``escape_with``."""
    return "".join((escape_with + c if c in to_escape else c) for c in s)


def get_file_ext(filename: str) -> str:
    """Returns the lowercase extension part of filename, without the dot."""
    pos = filename.rfind(".")
    if pos > -1:
        return filename[pos + 1 :].lower()
    else:
        return ""


def rem_file_ext(filename: str) -> str:
    """Returns the filename without extension."""
    pos = filename.rfind(".")
    if pos > -1:
        return filename[:pos]
    else:
        return filename


# TODO type hint number
def pluralize(number, word: str, decimals: int = 0, plural_word: Union[str, None] = None) -> str:
    """Returns a pluralized string with ``number`` in front of ``word``.

    Adds a 's' to s if ``number`` > 1.
    ``number``: The number to go in front of s
    ``word``: The word to go after number
    ``decimals``: The number of digits after the dot
    ``plural_word``: If the plural rule for word is more complex than adding a 's', specify a plural
    """
    number = round(number, decimals)
    plural_format = "%%1.%df %%s" % decimals
    if number > 1:
        if plural_word is None:
            word += "s"
        else:
            word = plural_word
    return plural_format % (number, word)


def format_time(seconds: int, with_hours: bool = True) -> str:
    """Transforms seconds in a hh:mm:ss string.

    If ``with_hours`` if false, the format is mm:ss.
    """
    minus = seconds < 0
    if minus:
        seconds *= -1
    m, s = divmod(seconds, 60)
    if with_hours:
        h, m = divmod(m, 60)
        r = "%02d:%02d:%02d" % (h, m, s)
    else:
        r = "%02d:%02d" % (m, s)
    if minus:
        return "-" + r
    else:
        return r


def format_time_decimal(seconds: int) -> str:
    """Transforms seconds in a strings like '3.4 minutes'."""
    minus = seconds < 0
    if minus:
        seconds *= -1
    if seconds < 60:
        r = pluralize(seconds, "second", 1)
    elif seconds < 3600:
        r = pluralize(seconds / 60.0, "minute", 1)
    elif seconds < 86400:
        r = pluralize(seconds / 3600.0, "hour", 1)
    else:
        r = pluralize(seconds / 86400.0, "day", 1)
    if minus:
        return "-" + r
    else:
        return r


SIZE_DESC = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
SIZE_VALS = tuple(1024**i for i in range(1, 9))


def format_size(size: int, decimal: int = 0, forcepower: int = -1, showdesc: bool = True) -> str:
    """Transform a byte count in a formatted string (KB, MB etc..).

    ``size`` is the number of bytes to format.
    ``decimal`` is the number digits after the dot.
    ``forcepower`` is the desired suffix. 0 is B, 1 is KB, 2 is MB etc.. if kept at -1, the suffix
    will be automatically chosen (so the resulting number is always below 1024).
    if ``showdesc`` is ``True``, the suffix will be shown after the number.
    Usage example::

        >>> format_size(1234, decimal=2, showdesc=True)
        '1.21 KB'
    """
    if forcepower < 0:
        i = 0
        while size >= SIZE_VALS[i]:
            i += 1
    else:
        i = forcepower
    if i > 0:
        div = SIZE_VALS[i - 1]
    else:
        div = 1
    size_format = "%%%d.%df" % (decimal, decimal)
    negative = size < 0
    divided_size = (0.0 + abs(size)) / div
    if decimal == 0:
        divided_size = ceil(divided_size)
    else:
        divided_size = ceil(divided_size * (10**decimal)) / (10**decimal)
    if negative:
        divided_size *= -1
    result = size_format % divided_size
    if showdesc:
        result += " " + SIZE_DESC[i]
    return result


def multi_replace(s: str, replace_from: Union[str, List[str]], replace_to: Union[str, List[str]] = "") -> str:
    """A function like str.replace() with multiple replacements.

    ``replace_from`` is a list of things you want to replace. Ex: ['a','bc','d']
    ``replace_to`` is a list of what you want to replace to.
    If ``replace_to`` is a list and has the same length as ``replace_from``, ``replace_from``
    items will be translated to corresponding ``replace_to``. A ``replace_to`` list must
    have the same length as ``replace_from``
    If ``replace_to`` is a string, all ``replace_from`` occurence will be replaced
    by that string.
    ``replace_from`` can also be a str. If it is, every char in it will be translated
    as if ``replace_from`` would be a list of chars. If ``replace_to`` is a str and has
    the same length as ``replace_from``, it will be transformed into a list.
    """
    if isinstance(replace_to, str) and (len(replace_from) != len(replace_to)):
        replace_to = [replace_to for _ in replace_from]
    if len(replace_from) != len(replace_to):
        raise ValueError("len(replace_from) must be equal to len(replace_to)")
    replace = list(zip(replace_from, replace_to))
    for r_from, r_to in [r for r in replace if r[0] in s]:
        s = s.replace(r_from, r_to)
    return s


# --- Files related


@log_io_error
@pathify
def delete_if_empty(path: Path, files_to_delete: List[str] = []) -> bool:
    """Deletes the directory at 'path' if it is empty or if it only contains files_to_delete."""
    if not path.exists() or not path.is_dir():
        return False
    contents = list(path.glob("*"))
    if any(p for p in contents if (p.name not in files_to_delete) or p.is_dir()):
        return False
    for p in contents:
        p.unlink()
    path.rmdir()
    return True


def open_if_filename(
    infile: Union[Path, str, IO],
    mode: str = "rb",
) -> Tuple[IO, bool]:
    """If ``infile`` is a string, it opens and returns it. If it's already a file object, it simply returns it.

    This function returns ``(file, should_close_flag)``. The should_close_flag is True is a file has
    effectively been opened (if we already pass a file object, we assume that the responsibility for
    closing the file has already been taken). Example usage::

        fp, shouldclose = open_if_filename(infile)
        dostuff()
        if shouldclose:
            fp.close()
    """
    if isinstance(infile, Path):
        return (infile.open(mode), True)
    if isinstance(infile, str):
        return (open(infile, mode), True)
    else:
        return (infile, False)


class FileOrPath:
    """Does the same as :func:`open_if_filename`, but it can be used with a ``with`` statement.

    Example::

        with FileOrPath(infile):
            dostuff()
    """

    def __init__(self, file_or_path: Union[Path, str], mode: str = "rb") -> None:
        self.file_or_path = file_or_path
        self.mode = mode
        self.mustclose = False
        self.fp: Union[IO, None] = None

    def __enter__(self) -> IO:
        self.fp, self.mustclose = open_if_filename(self.file_or_path, self.mode)
        return self.fp

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        if self.fp and self.mustclose:
            self.fp.close()