monkey/chaos_monkey/monkeyfs.py

from io import BytesIO
import os

__author__ = 'hoffer'

MONKEYFS_PREFIX = 'monkeyfs://'

open_orig = open

class VirtualFile(BytesIO):
    _vfs = {} #virtual File-System 

    def __init__(self, name, mode = 'r', buffering = None):
        if not name.startswith(MONKEYFS_PREFIX):
            name = MONKEYFS_PREFIX + name
        self.name = name
        self._mode = mode
        if VirtualFile._vfs.has_key(name):
            super(VirtualFile, self).__init__(self._vfs[name])
        else:
            super(VirtualFile, self).__init__('')

    def flush(self):
        super(VirtualFile, self).flush()
        VirtualFile._vfs[self.name] = self.getvalue()

    @staticmethod
    def getsize(path):
        return len(VirtualFile._vfs[path])

    @staticmethod
    def isfile(path):
        return VirtualFile._vfs.has_key(path)

def getsize(path):
    if path.startswith(MONKEYFS_PREFIX):
        return VirtualFile.getsize(path)
    else:
        return os.stat(path).st_size

def isfile(path):
    if path.startswith(MONKEYFS_PREFIX):
        return VirtualFile.isfile(path)
    else:
        return os.path.isfile(path)

def virtual_path(name):
    return "%s%s" % (MONKEYFS_PREFIX, name)

def open(name, mode='r', buffering=-1):
    #use normal open for regular paths, and our "virtual" open for monkeyfs:// paths
    if name.startswith(MONKEYFS_PREFIX):
        return VirtualFile(name, mode, buffering)
    else:
        return open_orig(name, mode=mode, buffering=buffering)