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
|
