File size: 4,389 Bytes
d1ed09d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
Disk management utilities.
"""

# Authors: Gael Varoquaux <gael dot varoquaux at normalesup dot org>
#          Lars Buitinck
# Copyright (c) 2010 Gael Varoquaux
# License: BSD Style, 3 clauses.


import os
import sys
import time
import errno
import shutil

from multiprocessing import util


try:
    WindowsError
except NameError:
    WindowsError = OSError


def disk_used(path):
    """ Return the disk usage in a directory."""
    size = 0
    for file in os.listdir(path) + ['.']:
        stat = os.stat(os.path.join(path, file))
        if hasattr(stat, 'st_blocks'):
            size += stat.st_blocks * 512
        else:
            # on some platform st_blocks is not available (e.g., Windows)
            # approximate by rounding to next multiple of 512
            size += (stat.st_size // 512 + 1) * 512
    # We need to convert to int to avoid having longs on some systems (we
    # don't want longs to avoid problems we SQLite)
    return int(size / 1024.)


def memstr_to_bytes(text):
    """ Convert a memory text to its value in bytes.
    """
    kilo = 1024
    units = dict(K=kilo, M=kilo ** 2, G=kilo ** 3)
    try:
        size = int(units[text[-1]] * float(text[:-1]))
    except (KeyError, ValueError) as e:
        raise ValueError(
            "Invalid literal for size give: %s (type %s) should be "
            "alike '10G', '500M', '50K'." % (text, type(text))) from e
    return size


def mkdirp(d):
    """Ensure directory d exists (like mkdir -p on Unix)
    No guarantee that the directory is writable.
    """
    try:
        os.makedirs(d)
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise


# if a rmtree operation fails in rm_subdirs, wait for this much time (in secs),
# then retry up to RM_SUBDIRS_N_RETRY times. If it still fails, raise the
# exception. this mechanism ensures that the sub-process gc have the time to
# collect and close the memmaps before we fail.
RM_SUBDIRS_RETRY_TIME = 0.1
RM_SUBDIRS_N_RETRY = 10


def rm_subdirs(path, onerror=None):
    """Remove all subdirectories in this path.

    The directory indicated by `path` is left in place, and its subdirectories
    are erased.

    If onerror is set, it is called to handle the error with arguments (func,
    path, exc_info) where func is os.listdir, os.remove, or os.rmdir;
    path is the argument to that function that caused it to fail; and
    exc_info is a tuple returned by sys.exc_info().  If onerror is None,
    an exception is raised.
    """

    # NOTE this code is adapted from the one in shutil.rmtree, and is
    # just as fast

    names = []
    try:
        names = os.listdir(path)
    except os.error:
        if onerror is not None:
            onerror(os.listdir, path, sys.exc_info())
        else:
            raise

    for name in names:
        fullname = os.path.join(path, name)
        delete_folder(fullname, onerror=onerror)


def delete_folder(folder_path, onerror=None, allow_non_empty=True):
    """Utility function to cleanup a temporary folder if it still exists."""
    if os.path.isdir(folder_path):
        if onerror is not None:
            shutil.rmtree(folder_path, False, onerror)
        else:
            # allow the rmtree to fail once, wait and re-try.
            # if the error is raised again, fail
            err_count = 0
            while True:
                files = os.listdir(folder_path)
                try:
                    if len(files) == 0 or allow_non_empty:
                        shutil.rmtree(
                            folder_path, ignore_errors=False, onerror=None
                        )
                        util.debug(
                            "Successfully deleted {}".format(folder_path))
                        break
                    else:
                        raise OSError(
                            "Expected empty folder {} but got {} "
                            "files.".format(folder_path, len(files))
                        )
                except (OSError, WindowsError):
                    err_count += 1
                    if err_count > RM_SUBDIRS_N_RETRY:
                        # the folder cannot be deleted right now. It maybe
                        # because some temporary files have not been deleted
                        # yet.
                        raise
                time.sleep(RM_SUBDIRS_RETRY_TIME)