Bug 1130336 - Create a function to determine: total, used and free disk space. r=jlund
new file mode 100644
--- /dev/null
+++ b/mozharness/base/diskutils.py
@@ -0,0 +1,156 @@
+"""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 os
+import sys
+import logging
+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]
+ 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, unicode) 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
--- a/mozharness/base/log.py
+++ b/mozharness/base/log.py
@@ -21,16 +21,29 @@ import traceback
FATAL_LEVEL = logging.CRITICAL + 10
logging.addLevelName(FATAL_LEVEL, 'FATAL')
# mozharness log levels.
DEBUG, INFO, WARNING, ERROR, CRITICAL, FATAL, IGNORE = (
'debug', 'info', 'warning', 'error', 'critical', 'fatal', 'ignore')
+LOG_LEVELS = {
+ DEBUG: logging.DEBUG,
+ INFO: logging.INFO,
+ WARNING: logging.WARNING,
+ ERROR: logging.ERROR,
+ CRITICAL: logging.CRITICAL,
+ FATAL: FATAL_LEVEL
+}
+
+# mozharness root logger
+ROOT_LOGGER = logging.getLogger()
+
+
# LogMixin {{{1
class LogMixin(object):
"""This is a mixin for any object to access similar logging
functionality -- more so, of course, for those objects with
self.config and self.log_obj, of course.
"""
def _log_level_at_least(self, level):
@@ -201,24 +214,17 @@ pre-context-line setting in error_list.)
# BaseLogger {{{1
class BaseLogger(object):
"""Create a base logging class.
TODO: status? There may be a status object or status capability in
either logging or config that allows you to count the number of
error,critical,fatal messages for us to count up at the end (aiming
for 0).
"""
- LEVELS = {
- DEBUG: logging.DEBUG,
- INFO: logging.INFO,
- WARNING: logging.WARNING,
- ERROR: logging.ERROR,
- CRITICAL: logging.CRITICAL,
- FATAL: FATAL_LEVEL
- }
+ LEVELS = LOG_LEVELS
def __init__(
self, log_level=INFO,
log_format='%(message)s',
log_date_format='%H:%M:%S',
log_name='test',
log_to_console=True,
log_dir='.',
@@ -252,17 +258,17 @@ class BaseLogger(object):
os.makedirs(self.log_dir)
self.abs_log_dir = os.path.abspath(self.log_dir)
def init_message(self, name=None):
if not name:
name = self.__class__.__name__
self.log_message("%s online at %s in %s" %
(name, datetime.now().strftime("%Y%m%d %H:%M:%S"),
- os.getcwd()))
+ os.getcwd()))
def get_logger_level(self, level=None):
if not level:
level = self.log_level
return self.LEVELS.get(level, logging.NOTSET)
def get_log_formatter(self, log_format=None, date_format=None):
if not log_format:
@@ -270,17 +276,17 @@ class BaseLogger(object):
if not date_format:
date_format = self.log_date_format
return logging.Formatter(log_format, date_format)
def new_logger(self, logger_name):
"""Create a new logger.
By default there are no handlers.
"""
- self.logger = logging.getLogger(logger_name)
+ self.logger = ROOT_LOGGER
self.logger.setLevel(self.get_logger_level())
self._clear_handlers()
if self.log_to_console:
self.add_console_handler()
if self.log_to_raw:
self.log_files['raw'] = '%s_raw.log' % self.log_name
self.add_file_handler(os.path.join(self.abs_log_dir,
self.log_files['raw']),
@@ -383,11 +389,18 @@ class MultiFileLogger(BaseLogger):
if self.get_logger_level(level) >= min_logger_level:
self.log_files[level] = '%s_%s.log' % (self.log_name,
level)
self.add_file_handler(os.path.join(self.abs_log_dir,
self.log_files[level]),
log_level=level)
+def numeric_log_level(level):
+ """Converts a mozharness log level (string) to the corresponding logger
+ level (number). This function makes possible to set the log level
+ in functions that do not inherit from LogMixin
+ """
+ return LOG_LEVELS[level]
+
# __main__ {{{1
if __name__ == '__main__':
pass
--- a/mozharness/base/script.py
+++ b/mozharness/base/script.py
@@ -1357,18 +1357,16 @@ class BaseScript(ScriptMixin, LogMixin,
"log_name": 'log',
"log_dir": log_dir,
"log_level": default_log_level,
"log_format": '%(asctime)s %(levelname)8s - %(message)s',
"log_to_console": True,
"append_to_log": False,
}
log_type = self.config.get("log_type", "multi")
- if log_type == "multi":
- log_config['logger_name'] = 'Multi'
for key in log_config.keys():
value = self.config.get(key, None)
if value is not None:
log_config[key] = value
if log_type == "multi":
self.log_obj = MultiFileLogger(**log_config)
else:
self.log_obj = SimpleFileLogger(**log_config)
new file mode 100644
--- /dev/null
+++ b/test/test_base_diskutils.py
@@ -0,0 +1,84 @@
+import mock
+import unittest
+from mozharness.base.diskutils import convert_to, DiskutilsError, DiskSize, DiskInfo
+
+
+class TestDiskutils(unittest.TestCase):
+ def test_convert_to(self):
+ # 0 is 0 regardless from_unit/to_unit
+ self.assertTrue(convert_to(size=0, from_unit='GB', to_unit='MB') == 0)
+ size = 524288 # 512 * 1024
+ # converting from/to same unit
+ self.assertTrue(convert_to(size=size, from_unit='MB', to_unit='MB') == size)
+
+ self.assertTrue(convert_to(size=size, from_unit='MB', to_unit='GB') == 512)
+
+ self.assertRaises(DiskutilsError,
+ lambda: convert_to(size='a string', from_unit='MB', to_unit='MB'))
+ self.assertRaises(DiskutilsError,
+ lambda: convert_to(size=0, from_unit='foo', to_unit='MB'))
+ self.assertRaises(DiskutilsError,
+ lambda: convert_to(size=0, from_unit='MB', to_unit='foo'))
+
+
+class TestDiskInfo(unittest.TestCase):
+
+ def testDiskinfo_to(self):
+ di = DiskInfo()
+ self.assertTrue(di.unit == 'bytes')
+ self.assertTrue(di.free == 0)
+ self.assertTrue(di.used == 0)
+ self.assertTrue(di.total == 0)
+ # convert to GB
+ di._to('GB')
+ self.assertTrue(di.unit == 'GB')
+ self.assertTrue(di.free == 0)
+ self.assertTrue(di.used == 0)
+ self.assertTrue(di.total == 0)
+
+
+class MockStatvfs(object):
+ def __init__(self):
+ self.f_bsize = 0
+ self.f_frsize = 0
+ self.f_blocks = 0
+ self.f_bfree = 0
+ self.f_bavail = 0
+ self.f_files = 0
+ self.f_ffree = 0
+ self.f_favail = 0
+ self.f_flag = 0
+ self.f_namemax = 0
+
+
+class TestDiskSpace(unittest.TestCase):
+
+ @mock.patch('mozharness.base.diskutils.os')
+ def testDiskSpacePosix(self, mock_os):
+ ds = MockStatvfs()
+ mock_os.statvfs.return_value = ds
+ di = DiskSize()._posix_size('/')
+ self.assertTrue(di.unit == 'bytes')
+ self.assertTrue(di.free == 0)
+ self.assertTrue(di.used == 0)
+ self.assertTrue(di.total == 0)
+
+ @mock.patch('mozharness.base.diskutils.ctypes')
+ def testDiskSpaceWindows(self, mock_ctypes):
+ mock_ctypes.windll.kernel32.GetDiskFreeSpaceExA.return_value = 0
+ mock_ctypes.windll.kernel32.GetDiskFreeSpaceExW.return_value = 0
+ di = DiskSize()._windows_size('/c/')
+ self.assertTrue(di.unit == 'bytes')
+ self.assertTrue(di.free == 0)
+ self.assertTrue(di.used == 0)
+ self.assertTrue(di.total == 0)
+
+ @mock.patch('mozharness.base.diskutils.os')
+ @mock.patch('mozharness.base.diskutils.ctypes')
+ def testUnspportedPlafrom(self, mock_ctypes, mock_os):
+ mock_os.statvfs.side_effect = AttributeError('')
+ self.assertRaises(AttributeError, lambda: DiskSize()._posix_size('/'))
+ mock_ctypes.windll.kernel32.GetDiskFreeSpaceExW.side_effect = AttributeError('')
+ mock_ctypes.windll.kernel32.GetDiskFreeSpaceExA.side_effect = AttributeError('')
+ self.assertRaises(AttributeError, lambda: DiskSize()._windows_size('/'))
+ self.assertRaises(DiskutilsError, lambda: DiskSize().get_size(path='/', unit='GB'))