blender-benchmark-bundle/benchmark/foundation/util.py

from . import progress
from .third_party import dateutil
import os
import requests
import tarfile
import zipfile


def humanReadableTimeDifference(seconds):
    """
    Convert time difference in seconds to a human readable format.

    For example, time difference of 125 seconds will be returned as 2:05
    """

    hours = int(seconds) // 60 // 60
    seconds = seconds - hours * 60 * 60
    minutes = int(seconds) // 60
    seconds = seconds - minutes * 60
    if hours == 0:
        return "%02d:%05.2f" % (minutes, seconds)
    else:
        return "%02d:%02d:%05.2f" % (hours, minutes, seconds)


def humanReadableTimeToSeconds(time):
    """
    Convert human readable string like HH:MM:SS to seconds.
    """

    tokens = time.split(".")
    result = 0
    if len(tokens) == 2:
        result = float("0." + tokens[1])
    mult = 1
    for token in reversed(tokens[0].split(":")):
        result += int(token) * mult
        mult *= 60
    return result


def queryMainScene(filepath, callbacks):
    """
    Return the equivalent to bpy.context.scene
    """

    from .blendfile import blendfile

    with blendfile.open_blend(filepath) as blend:
        # There is no bpy.context.scene, we get it from the main window
        window_manager = [block for block in blend.blocks
                          if block.code == b'WM'][0]
        window = window_manager.get_pointer(b'winactive')
        screen = window.get_pointer(b'screen')
        scene = screen.get_pointer(b'scene')

        output = []
        for callback in callbacks:
            output.append(callback(scene))
        return output


def queryCurrentFrame(filepath):
    """
    Get frame number to render.
    """

    def get_cfra(scene):
        return scene.get((b'r', b'cfra'))
    cfra, = queryMainScene(filepath, [get_cfra])
    return cfra


def humanReadableSizeToMegabytes(size):
    if size[-1] == 'K':
        return float(size[:-1]) / 1024
    elif size[-1] == 'M':
        return float(size[:-1])
    else:
        return float(size)


def humanReadableSize(size):
    return "{} Mb" . format(size)


def downloadFile(url, filename):
    """
    Download file form given UTR and save it to filename
    """
    r = requests.get(url, stream=True)
    downloaded_size = 0
    total_size = 0
    if 'Content-length' in r.headers:
        total_size = r.headers['Content-length']
    with open(filename, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk:
                f.write(chunk)
                downloaded_size += len(chunk)
                if total_size != 0:
                    progress.progress(downloaded_size, total_size)
    if total_size != 0:
        progress.progressClear()


def unpackArchive(filename, directory):
    """
    Unpack archive <filename> to given <directory>
    """
    if filename.endswith(".tar.bz2"):
        tar = tarfile.open(name=filename, mode="r:bz2")
        tar.extractall(directory)
    elif filename.endswith(".zip"):
        zip_ref = zipfile.ZipFile(filename, 'r')
        zip_ref.extractall(directory)
        zip_ref.close()
    else:
        # TODO(sergey): Need to support more archive types.
        pass


def stripSensitiveInfo(s):
    """
    Strip any possibly sensitive information we want to avoid from the logs.
    """
    root_dir = getBundleRootDirectory()
    s = s.replace(root_dir, "<FARM_DIR>")
    return s


def blenderCommitUnixTimestamp(commit_date, commit_time):
    """
    Convert commit build time and time to unix timestamp
    """
    date_time = commit_date + " " + commit_time
    return dateutil.parser.parse(date_time)


########################################
#  Directories manipulation.
########################################


def getBundleRootDirectory():
    """
    Get fuill file path to the root directory of benchmark bundle.
    """
    script_directory = os.path.dirname(os.path.realpath(__file__))
    return os.path.dirname(os.path.dirname(script_directory))


def getGlobalConfigDirectory():
    """
    Get configuration directory global for the all components of this bundle.
    """
    return os.path.join(getBundleRootDirectory(), "config")


########################################
#  Devine name manipulation.
########################################


def deviceInfoAsString(device_info):
    """
    Convert device information to a single string.
    """
    device_type = device_info['device_type']
    compute_devices = device_info['compute_devices']
    num_devices = len(compute_devices)
    if num_devices == 1:
        return compute_devices[0].replace(" (Display)", "")
    elif num_devices == 2:
        if compute_devices[0] == compute_devices[1]:
            return compute_devices[0] + " DUO"
        raise Exception("Needs implementation")

########################################
#  Graphs manipulation.
########################################


def generateBarColor(index, alpha=None):
    """
    Generate unique looking color for a given bar index.
    """
    builtin_colors = ((255, 99, 132),
                      (255, 159, 64),
                      (255, 205, 86),
                      (75, 192, 192),
                      (54, 162, 235),
                      (153, 102, 255),
                      (201, 203, 207),
                      (48, 103, 204),
                      (220, 56, 18),
                      (254, 155, 0),
                      (15, 147, 25))
    color = (0, 0, 0)
    if index >= 0 and index < len(builtin_colors):
        color = builtin_colors[index]
    if alpha is None:
        return "rgb({}, {}, {})" . format(
            str(color[0]), str(color[1]), str(color[2]))
    else:
        return "rgba({}, {}, {}, {})" . format(
            str(color[0]), str(color[1]), str(color[2]), str(alpha))