summaryrefslogtreecommitdiffstats
path: root/testing/mozharness/mozharness/base/diskutils.py
blob: 757a6ffb7ab260951cadaae5655660c61cca201d (plain)
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
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""Disk utility module, no mixins here!

    examples:
    1) get disk size
    from mozharness.base.diskutils import DiskInfo, DiskutilsError
    ...
    try:
        DiskSize().get_size(path='/', unit='Mb')
    except DiskutilsError as e:
        # manage the exception e.g: log.error(e)
        pass
    log.info("%s" % di)


    2) convert disk size:
    from mozharness.base.diskutils import DiskutilsError, convert_to
    ...
    file_size = <function that gets file size in bytes>
    # convert file_size to GB
    try:
        file_size = convert_to(file_size, from_unit='bytes', to_unit='GB')
    except DiskutilsError as e:
        # manage the exception e.g: log.error(e)
        pass

"""
import ctypes
import logging
import os
import sys

from six import string_types

from mozharness.base.log import INFO, numeric_log_level

# use mozharness log
log = logging.getLogger(__name__)


class DiskutilsError(Exception):
    """Exception thrown by Diskutils module"""

    pass


def convert_to(size, from_unit, to_unit):
    """Helper method to convert filesystem sizes to kB/ MB/ GB/ TB/
    valid values for source_format and destination format are:
        * bytes
        * kB
        * MB
        * GB
        * TB
     returns: size converted from source_format to destination_format.
    """
    sizes = {
        "bytes": 1,
        "kB": 1024,
        "MB": 1024 * 1024,
        "GB": 1024 * 1024 * 1024,
        "TB": 1024 * 1024 * 1024 * 1024,
    }
    try:
        df = sizes[to_unit]
        sf = sizes[from_unit]
        # pylint --py3k W1619
        return size * sf / df
    except KeyError:
        raise DiskutilsError("conversion error: Invalid source or destination format")
    except TypeError:
        raise DiskutilsError("conversion error: size (%s) is not a number" % size)


class DiskInfo(object):
    """Stores basic information about the disk"""

    def __init__(self):
        self.unit = "bytes"
        self.free = 0
        self.used = 0
        self.total = 0

    def __str__(self):
        string = ["Disk space info (in %s)" % self.unit]
        string += ["total: %s" % self.total]
        string += ["used: %s" % self.used]
        string += ["free: %s" % self.free]
        return " ".join(string)

    def _to(self, unit):
        from_unit = self.unit
        to_unit = unit
        self.free = convert_to(self.free, from_unit=from_unit, to_unit=to_unit)
        self.used = convert_to(self.used, from_unit=from_unit, to_unit=to_unit)
        self.total = convert_to(self.total, from_unit=from_unit, to_unit=to_unit)
        self.unit = unit


class DiskSize(object):
    """DiskSize object"""

    @staticmethod
    def _posix_size(path):
        """returns the disk size in bytes
        disk size is relative to path
        """
        # we are on a POSIX system
        st = os.statvfs(path)
        disk_info = DiskInfo()
        disk_info.free = st.f_bavail * st.f_frsize
        disk_info.used = (st.f_blocks - st.f_bfree) * st.f_frsize
        disk_info.total = st.f_blocks * st.f_frsize
        return disk_info

    @staticmethod
    def _windows_size(path):
        """returns size in bytes, works only on windows platforms"""
        # we're on a non POSIX system (windows)
        # DLL call
        disk_info = DiskInfo()
        dummy = ctypes.c_ulonglong()  # needed by the dll call but not used
        total = ctypes.c_ulonglong()  # stores the total space value
        free = ctypes.c_ulonglong()  # stores the free space value
        # depending on path format (unicode or not) and python version (2 or 3)
        # we need to call GetDiskFreeSpaceExW or GetDiskFreeSpaceExA
        called_function = ctypes.windll.kernel32.GetDiskFreeSpaceExA
        if isinstance(path, string_types) or sys.version_info >= (3,):
            called_function = ctypes.windll.kernel32.GetDiskFreeSpaceExW
        # we're ready for the dll call. On error it returns 0
        if (
            called_function(
                path, ctypes.byref(dummy), ctypes.byref(total), ctypes.byref(free)
            )
            != 0
        ):
            # success, we can use the values returned by the dll call
            disk_info.free = free.value
            disk_info.total = total.value
            disk_info.used = total.value - free.value
        return disk_info

    @staticmethod
    def get_size(path, unit, log_level=INFO):
        """Disk info stats:
              total => size of the disk
              used  => space used
              free  => free space
        In case of error raises a DiskutilError Exception
        """
        try:
            # let's try to get the disk size using os module
            disk_info = DiskSize()._posix_size(path)
        except AttributeError:
            try:
                # os module failed. let's try to get the size using
                # ctypes.windll...
                disk_info = DiskSize()._windows_size(path)
            except AttributeError:
                # No luck! This is not a posix nor window platform
                # raise an exception
                raise DiskutilsError("Unsupported platform")

        disk_info._to(unit)
        lvl = numeric_log_level(log_level)
        log.log(lvl, msg="%s" % disk_info)
        return disk_info