marmots/tpers.py

import cPickle
import pdb
import traceback

def Decorator(dg):
    def DecoratorNew(*rgarg, **mparg):  
        def CallDecorator(o):
            return dg(o, *rgarg, **mparg) or o
        return CallDecorator
    return DecoratorNew

@Decorator
def Version(cls, vr):
    cls._version = vr

def RenameFrom(stName, namespace):
    def SetOldName(cls):
        namespace[stName] = cls
        return cls
    return SetOldName

class TransientSetter(object):
    def __init__(self, obj):
        self.obj = obj
        try:
            self._fWriteOld = obj._fWriteToPersistent
        except Exception:
            self._fWriteOld = True
    def __enter__(self):
        self.obj._fWriteToPersistent = False
        return self.obj
    def __exit__(self, type, value, traceback):
        self.obj._fWriteToPersistent = self._fWriteOld

@Version(1)
class TPrs(object):
    def __init__(self, *rgarg, **mparg):
        self._persistent = TPM()
        self._fWriteToPersistent = True
        self._versionCurr = self._version 
        self._fPersistent = True
        self.InitPersistent(*rgarg, **mparg)
        with self.SetTransiently():
            self.InitTransient()

    def InitPersistent(self, *rgarg, **mparg):
        pass

    def InitTransient(self):
        pass

    def SetTransiently(self):
        return TransientSetter(self)

    def MakeTransient(self, key, value):
        try:
            del self._persistent[key]
        except Exception:
            pass
        with self.SetTransiently():
            setattr(self, key, value)

    def FPersist(self):
        return self._fPersistent

    def DontPersist(self):
        self._fPersistent = False

    def UpgradeFrom(self, versionOld):
        pass

    def __getattr__(self, key):
        stPersistent = "no persistent yet"
        if key != "_persistent":
            try:
                return self._persistent[key]
            except Exception:
                stPersistent = str(self._persistent)
        raise AttributeError("Attribute '" + key + "' not found in object with type '" + type(self).__name__ + "':" + stPersistent + " __dict__:" + str(self.__dict__))

    __obj_setattr = object.__setattr__
    __obj_delattr = object.__delattr__
    def __setattr__(self, key, value):
        prop = getattr(self.__class__, key, None)
        if isinstance(prop, property) and prop.fset:
            prop.fset(self, value)
            return
        try:
            if self._fWriteToPersistent or key in self._persistent:
                self._persistent[key] = value
                try:
                    self.__obj_delattr(key)
                except AttributeError:
                    pass
                return
        except Exception: pass
        self.__obj_setattr(key, value)

    def __delattr__(self, key):
        try:
            del self._persistent[key]
        except Exception:
            self.__obj_delattr(key)

    def __reduce__(self):
        if self.FPersist():
            return (NewObjWithClass, (type(self),), self._persistent)
        else:
            return (NewObjWithClass, (None,))

    def __setstate__(self, state):
        self._persistent = state
        self._fWriteToPersistent = True
        if self._versionCurr != self._version:
            Odb.rgtprsToUpgrade.append(self)
        Odb.rgtprsToInit.append(self)
    def _Upgrade(self):
        self.UpgradeFrom(self._versionCurr)
        self._versionCurr = self._version

class Odb(object):
    @staticmethod
    def Save(tprs, fn):
        with open(fn, "w") as fl:
            cPickle.dump(tprs, fl)

    @classmethod
    def Load(cls, fn):
        cls.rgtprsToUpgrade = []
        cls.rgtprsToInit = []
        try:
            with open(fn, "r") as fl:
                tprs = cPickle.load(fl)
        except Exception as e:
            print "error unpickling:", e
            traceback.print_exc()
            return None
        for tprsToUpgrade in cls.rgtprsToUpgrade:
            tprsToUpgrade._Upgrade()
        for tprsToInit in cls.rgtprsToInit:
            with tprsToInit.SetTransiently():
                tprsToInit.InitTransient()
        if len(cls.rgtprsToUpgrade) > 0:
            fnUpgraded = fn + ".upgraded"
            Odb.Save(tprs, fnUpgraded)
            tprs = Odb.Load(fnUpgraded)
        cls.rgtprsToUpgrade = None
        cls.rgtprsToInit = None
        return tprs

clsList = list
clsDict = dict
clsSet = set
clsTuple = tuple

def MkTransparent(value):
    if type(value) == list:
        return TPL(value)
    elif type(value) == dict:
        return TPM(value)
    elif type(value) == set:
        return TPS(value)
    elif type(value) == tuple:
        return TPT(value)
    else:
        return value

def RgTransparent(rg):
    if rg == None: return []
    return [MkTransparent(x) for x in rg]

def MpTransparent(mp):
    if isinstance(mp, TPM): return mp
    mpNew = {}
    if mp != None:
        for k,v in mp.iteritems():
            mpNew[k] = MkTransparent(v)
    return mpNew

def FPersist(obj):
    return (not isinstance(obj, TPrs)) or obj.FPersist()

def MkPersistable(obj):
    if FPersist(obj):
        return obj
    return None

def NewObjWithClass(cls):
    if cls == None:
        return None
    return cls.__new__(cls)

class TPL(clsList):
    __pl_setitem = clsList.__setitem__
    __pl_setslice = clsList.__setslice__
    __pl_iadd = clsList.__iadd__
    __pl_append = clsList.append
    __pl_insert = clsList.insert
    __pl_extend = clsList.extend

    def __init__(self, rg = None):
        clsList.__init__(self, RgTransparent(rg))

    def __setitem__(self, i, item):
        return self.__pl_setitem(i, MkTransparent(item))

    def __setslice__(self, i, j, other):
        return self.__pl_setslice(i, j, RgTransparent(other))

    def __iadd__(self, other):
        return self.__pl_iadd(RgTransparent(other))

    def append(self, item):
        return self.__pl_append(MkTransparent(item))

    def insert(self, i, item):
        return self.__pl_insert(i, MkTransparent(item))

    def extend(self, other):
        return self.__pl_extend(RgTransparent(other))

    def __reduce__(self):
        return (MkTransparent, ([MkPersistable(x) for x in self],))

# Strip weak references rather than replacing them with None
class TPLS(TPL):
    def __reduce__(self):
        return (MkTransparent, ([x for x in self if FPersist(x)],))

class TPM(clsDict):
    __pm_setitem = clsDict.__setitem__
    __pm_update = clsDict.update

    def __init__(self, *rgarg, **kwarg):
        clsDict.__init__(self, MpTransparent(dict(*rgarg, **kwarg)))

    def __setitem__(self, k, v):
        return self.__pm_setitem(k, MkTransparent(v))

    def update(self, other, **kwargs):
        return self.__pm_update(MpTransparent(other), **MpTransparent(kwargs))

    def __reduce__(self):
        mp = {}
        for k, v in self.iteritems():
            if FPersist(k):
                mp[k] = MkPersistable(v)
        return (MkTransparent, (mp,))

class TPS(clsSet):
    __ps_update = clsSet.update
    __ps_ior = clsSet.__ior__
    __ps_add = clsSet.add

    def __init__(self, rg = None):
        clsSet.__init__(self, RgTransparent(rg))

    def update(self, *rgother):
        return self.__ps_update(*RgTransparent(rgother))

    def __ior__(self, other):
        return self.__ps_ior(MkTransparent(other))

    def add(self, elem):
        return self.__ps_add(MkTransparent(elem))

    def __reduce__(self):
        return (MkTransparent, (set([x for x in self if FPersist(x)]),))

class TPT(clsTuple):
    def __new__(cls, tup = None):
        return clsTuple.__new__(cls, RgTransparent(tup))

    def __reduce__(self):
        return (MkTransparent, (tuple([MkPersistable(x) for x in self]),))