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