1"""Disk utility module, no mixins here! 2 3 examples: 4 1) get disk size 5 from mozharness.base.diskutils import DiskInfo, DiskutilsError 6 ... 7 try: 8 DiskSize().get_size(path='/', unit='Mb') 9 except DiskutilsError as e: 10 # manage the exception e.g: log.error(e) 11 pass 12 log.info("%s" % di) 13 14 15 2) convert disk size: 16 from mozharness.base.diskutils import DiskutilsError, convert_to 17 ... 18 file_size = <function that gets file size in bytes> 19 # convert file_size to GB 20 try: 21 file_size = convert_to(file_size, from_unit='bytes', to_unit='GB') 22 except DiskutilsError as e: 23 # manage the exception e.g: log.error(e) 24 pass 25 26""" 27import ctypes 28import os 29import sys 30import logging 31from mozharness.base.log import INFO, numeric_log_level 32 33# use mozharness log 34log = logging.getLogger(__name__) 35 36 37class DiskutilsError(Exception): 38 """Exception thrown by Diskutils module""" 39 pass 40 41 42def convert_to(size, from_unit, to_unit): 43 """Helper method to convert filesystem sizes to kB/ MB/ GB/ TB/ 44 valid values for source_format and destination format are: 45 * bytes 46 * kB 47 * MB 48 * GB 49 * TB 50 returns: size converted from source_format to destination_format. 51 """ 52 sizes = {'bytes': 1, 53 'kB': 1024, 54 'MB': 1024 * 1024, 55 'GB': 1024 * 1024 * 1024, 56 'TB': 1024 * 1024 * 1024 * 1024} 57 try: 58 df = sizes[to_unit] 59 sf = sizes[from_unit] 60 return size * sf / df 61 except KeyError: 62 raise DiskutilsError('conversion error: Invalid source or destination format') 63 except TypeError: 64 raise DiskutilsError('conversion error: size (%s) is not a number' % size) 65 66 67class DiskInfo(object): 68 """Stores basic information about the disk""" 69 def __init__(self): 70 self.unit = 'bytes' 71 self.free = 0 72 self.used = 0 73 self.total = 0 74 75 def __str__(self): 76 string = ['Disk space info (in %s)' % self.unit] 77 string += ['total: %s' % self.total] 78 string += ['used: %s' % self.used] 79 string += ['free: %s' % self.free] 80 return " ".join(string) 81 82 def _to(self, unit): 83 from_unit = self.unit 84 to_unit = unit 85 self.free = convert_to(self.free, from_unit=from_unit, to_unit=to_unit) 86 self.used = convert_to(self.used, from_unit=from_unit, to_unit=to_unit) 87 self.total = convert_to(self.total, from_unit=from_unit, to_unit=to_unit) 88 self.unit = unit 89 90 91class DiskSize(object): 92 """DiskSize object 93 """ 94 @staticmethod 95 def _posix_size(path): 96 """returns the disk size in bytes 97 disk size is relative to path 98 """ 99 # we are on a POSIX system 100 st = os.statvfs(path) 101 disk_info = DiskInfo() 102 disk_info.free = st.f_bavail * st.f_frsize 103 disk_info.used = (st.f_blocks - st.f_bfree) * st.f_frsize 104 disk_info.total = st.f_blocks * st.f_frsize 105 return disk_info 106 107 @staticmethod 108 def _windows_size(path): 109 """returns size in bytes, works only on windows platforms""" 110 # we're on a non POSIX system (windows) 111 # DLL call 112 disk_info = DiskInfo() 113 dummy = ctypes.c_ulonglong() # needed by the dll call but not used 114 total = ctypes.c_ulonglong() # stores the total space value 115 free = ctypes.c_ulonglong() # stores the free space value 116 # depending on path format (unicode or not) and python version (2 or 3) 117 # we need to call GetDiskFreeSpaceExW or GetDiskFreeSpaceExA 118 called_function = ctypes.windll.kernel32.GetDiskFreeSpaceExA 119 if isinstance(path, unicode) or sys.version_info >= (3,): 120 called_function = ctypes.windll.kernel32.GetDiskFreeSpaceExW 121 # we're ready for the dll call. On error it returns 0 122 if called_function(path, 123 ctypes.byref(dummy), 124 ctypes.byref(total), 125 ctypes.byref(free)) != 0: 126 # success, we can use the values returned by the dll call 127 disk_info.free = free.value 128 disk_info.total = total.value 129 disk_info.used = total.value - free.value 130 return disk_info 131 132 @staticmethod 133 def get_size(path, unit, log_level=INFO): 134 """Disk info stats: 135 total => size of the disk 136 used => space used 137 free => free space 138 In case of error raises a DiskutilError Exception 139 """ 140 try: 141 # let's try to get the disk size using os module 142 disk_info = DiskSize()._posix_size(path) 143 except AttributeError: 144 try: 145 # os module failed. let's try to get the size using 146 # ctypes.windll... 147 disk_info = DiskSize()._windows_size(path) 148 except AttributeError: 149 # No luck! This is not a posix nor window platform 150 # raise an exception 151 raise DiskutilsError('Unsupported platform') 152 153 disk_info._to(unit) 154 lvl = numeric_log_level(log_level) 155 log.log(lvl, msg="%s" % disk_info) 156 return disk_info 157