safe_wrap.py

import threading


class SafeWrap():
    """
    Wrap an unsafe sdk (like onedrivesdk, for example) in a thread lock.
    Derive from this if you want to translate exceptions, or change the
    type set
    """
    __safe_types = (str, int, float, type(None), bool, bytes, type(iter(bytes())), type(iter(str())))

    def __init__(self, o, lock=None):
        self.__o = o
        self.__lock = lock or threading.RLock()

    @classmethod
    def __safe_type(cls, typ):
        ok = isinstance(typ, cls.__safe_types)
        return ok

    def __wrap(self, ret):
        if self.__safe_type(ret):
            return ret
        return SafeWrap(ret, self.__lock)

    def __len__(self):
        return len(self.__o)

    def __iter__(self):
        return self.__wrap(iter(self.__o))

    def __next__(self):
        return self.__wrap(next(self.__o))

    def __getitem__(self, k):
        return self.__wrap(self.__o[k])

    def __getattr__(self, k):
        return self.__wrap(getattr(self.__o, k))

    def __call__(self, *a, **kw):
        try:
            with self.__lock:
                return self.__wrap(self.__o(*a, **kw))
        except Exception as e:
            new_e = self.__translate_exception(e)
            if new_e is not None:
                raise new_e
            raise

    def __translate_exception(self, e):
        ...

def test_wrap():
    class Foo():
        def __init__(self, x):
            self.name = x

        def api1(self):
            return b'hi'

        def api2(self):
            return Foo("subinstance")

        def api3(self):
            return {"x":Foo("subinstance")}


    cli = Foo("conn")

    w = SafeWrap(cli)

    assert type(w.api1) is SafeWrap


    assert w.api1() == b'hi'
    assert w.api2().name == "subinstance"
    assert type(w.api2()) == SafeWrap
    assert type(w.api3()["x"]) == SafeWrap
    assert [k for k in w.api3()] == ["x"]
    assert type([v for v in w.api3().values()][0]) == SafeWrap