File: util.py

package info (click to toggle)
dupeguru 4.3.1-6
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 3,604 kB
  • sloc: python: 16,846; ansic: 424; makefile: 123
file content (323 lines) | stat: -rw-r--r-- 10,435 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
# 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()