cache.py

import os
import pickle
import operator
from snimpy import mib, manager
from snimpy.smi import smi, ffi

_cache_file = '/tmp/smi_mib_cache'


class LazyLoad(object):
    pass


class LazyLoadType(LazyLoad):
    def __init__(self, module, name):
        self.module = module
        self.name = name

    def __call__(self):
        return getattr(
            operator.attrgetter(self.module.split('.', 1)[1])(
                __import__(self.module)),
            self.name)


class LazyLoadProxy(LazyLoad):
    def __init__(self, *args):
        self.args = args

    def __call__(self):
        args = (
            a() if isinstance(a, LazyLoad) else a
            for a in self.args
        )
        return ProxyNode(*args)


def load_cache():
    if not os.path.exists(_cache_file):
        return {}, {}
    with open(_cache_file, 'r') as fh:
        return pickle.load(fh)

def save_cache(obj):
    with open(_cache_file, 'w') as fh:
        return pickle.dump(obj, fh)


_name_cache, _mib_cache = load_cache()



class ProxyNode(object):
    def __init__(self, mib, name, type, node=None):
        self._mib = mib
        self._name = name
        self._type = type
        self._node = node

    @classmethod
    def _ensure_proxied(cls, obj):
        if isinstance(obj, mib.Node):
            mod = ffi.string(smi.smiGetNodeModule(obj.node).name)
            return ProxyNode(mod, str(obj), type(obj), node=obj)
        elif isinstance(obj, list):
            return [cls._ensure_proxied(i) for i in obj]
        return obj

    @classmethod
    def _ensure_lazy(cls, obj):
        if isinstance(obj, type):
            return LazyLoadType(obj.__module__, obj.__name__)
        elif isinstance(obj, ProxyNode):
            lazy_type = cls._ensure_lazy(obj._type)
            return LazyLoadProxy(obj._mib, obj._name, lazy_type)
        elif isinstance(obj, list):
            return [cls._ensure_lazy(i) for i in obj]
        else:
            return obj

    @classmethod
    def _ensure_not_lazy(cls, obj):
        if isinstance(obj, LazyLoad):
            return obj()
        elif isinstance(obj, list):
            return [cls._ensure_not_lazy(i) for i in obj]
        else:
            return obj

    def __getattr__(self, item):
        try:
            v = _mib_cache[self._mib][self._type.__name__][self._name][item]
        except KeyError:
            node = self._node
            if node is None:
                self._node = node = mib.get(self._mib, self._name)
                assert type(node) == self._type
            v = self._ensure_proxied(getattr(node, item))
            _mib_cache[self._mib][self._type.__name__][self._name][item] = self._ensure_lazy(v)

            save_cache((_name_cache, _mib_cache))
        else:
            v = self._ensure_not_lazy(v)

        setattr(self, item, v)
        return v



class Manager(manager.Manager):

    def _locate(self, attribute):
        for m, cache in _mib_cache.items():
            try:
                for cls in _lookups.keys():
                    if attribute in cache[cls.__name__]:
                        return m, ProxyNode(m, attribute, cls)
                else:
                    return m, mib.get(m, attribute)
            except mib.SMIException:
                pass
        raise AttributeError("{0} is not an attribute".format(attribute))

    def __getattribute__(self, attribute):
        if attribute.startswith("_"):
            return object.__getattribute__(self, attribute)
        m, a = self._locate(attribute)
        if isinstance(a, ProxyNode):
            attr_type = a._type
        else:
            attr_type = type(a)
        if issubclass(attr_type, mib.Scalar):
            oid, result = self._session.get(a.oid + (0,))[0]
            if result is not None:
                try:
                    return a.type(a, result)
                except ValueError:
                    if self._loose:
                        return result
                    raise
            return None
        elif issubclass(attr_type, mib.Column):
            return manager.ProxyColumn(self._session, a, self._loose)
        raise NotImplementedError

_lookups = {
    mib.Node: mib.getNodes,
    mib.Scalar: mib.getScalars,
    mib.Table: mib.getTables,
    mib.Column: mib.getColumns,
}

def load(mibname):
    try:
        m = _name_cache[mibname]
        cache_entry = _mib_cache[m]
    except KeyError:
        m = mib.load(mibname)
        cache_entry = {
            mib_obj_type.__name__: {str(o): {} for o in getter(m)}
            for mib_obj_type, getter in _lookups.items()
        }
        _name_cache[mibname] = m
        _mib_cache[m] = cache_entry
        save_cache((_name_cache, _mib_cache))

    for cls in (mib.Scalar, mib.Column):
        for o in cache_entry[cls.__name__].keys():
            setattr(Manager, o, 1)