X-Git-Url: https://git.openstreetmap.org./osqa.git/blobdiff_plain/07fe485fb0a119fb87f8ec948185b5acdd65300e..8b52676543e3f5f1350d8d0eeaedee6fe35f49b3:/forum/modules/decorators.py

diff --git a/forum/modules/decorators.py b/forum/modules/decorators.py
index a4c1fad..8fbee7e 100644
--- a/forum/modules/decorators.py
+++ b/forum/modules/decorators.py
@@ -1,59 +1,112 @@
 import inspect
 
 class DecoratableObject(object):
+    MODE_OVERRIDE = 0
+    MODE_PARAMS = 1
+    MODE_RESULT = 2
+
     def __init__(self, fn):
+        a = inspect.getargspec(fn)
         self._callable = fn
+        self._params_decoration = None
+        self._result_decoration = None
 
-    def decorate(self, fn, needs_origin):
+    def _decorate(self, fn, needs_origin, method=False):
         origin = self._callable
 
         if needs_origin:
-            self._callable = lambda *args, **kwargs: fn(origin, *args, **kwargs)
+            if method:
+                self._callable = lambda inst, *args, **kwargs: fn(inst, origin, *args, **kwargs)
+            else:
+                self._callable = lambda *args, **kwargs: fn(origin, *args, **kwargs)
         else:
-            self._callable = lambda *args, **kwargs: fn(*args, **kwargs)
+            self._callable = fn
 
-    def __call__(self, *args, **kwargs):
-        return self._callable(*args, **kwargs)
+    def _decorate_params(self, fn):
+        if not self._params_decoration:
+            self._params_decoration = []
 
+        self._params_decoration.append(fn)
 
-def decoratable(fn):
-    return DecoratableObject(fn)
+    def _decorate_result(self, fn):
+        if not self._result_decoration:
+            self._result_decoration = []
 
-def decoratable_method(fn):
-    obj = DecoratableObject(fn)
-    def decorated(self, *args, **kwargs):
-        return obj(self, *args, **kwargs)
+        self._result_decoration.append(fn)
 
-    decorated.__obj = obj
-    return decorated
+    def __call__(self, *args, **kwargs):
+        if self._params_decoration:
+            for dec in self._params_decoration:
+                args, kwargs = dec(*args, **kwargs)
 
-decoratable.method = decoratable_method
+        res = self._callable(*args, **kwargs)
 
-def decorate(origin, needs_origin=True):
-    if not isinstance(origin, DecoratableObject):
-        if hasattr(origin, '__obj'):
-            def decorator(fn):
-                origin.__obj.decorate(fn, needs_origin)
-                return origin
-            return decorator
+        if self._result_decoration:
+            for dec in self._result_decoration:
+                res = dec(res)
+
+        return res
 
-        raise TypeError('Not a decoratable function: %s' % origin.__name__)
 
+def _create_decorator(origin, needs_origin, mode, method=False):
     def decorator(fn):
-        origin.decorate(fn, needs_origin)
-        return origin
+        if mode == DecoratableObject.MODE_OVERRIDE:
+            origin._decorate(fn, needs_origin, method=method)
+        elif mode == DecoratableObject.MODE_PARAMS:
+            origin._decorate_params(fn)
+        elif mode == DecoratableObject.MODE_RESULT:
+            origin._decorate_result(fn)
+
+        return fn
 
     return decorator
 
 
-def decorate_all(module):
-    [setattr(module, n, decoratable(f)) for n, f in
-        [(n, getattr(module, n)) for n in dir(module)]
-        if (callable(f)) and (not inspect.isclass(f)) and (f.__module__ == module.__name__)]
+def _decorate_method(origin, needs_origin, mode):
+    if not hasattr(origin, '_decoratable_obj'):
+        name = origin.__name__
+        cls = origin.im_class
+
+        decoratable = DecoratableObject(origin)
+
+        def decoratable_method(self, *args, **kwargs):
+            return decoratable(self, *args, **kwargs)
+
+        decoratable_method._decoratable_obj = decoratable
+        setattr(cls, name, decoratable_method)
+    else:
+        decoratable = origin._decoratable_obj
 
+    return _create_decorator(decoratable, needs_origin, mode, method=True)
 
+def _decorate_function(origin, needs_origin, mode):
+    if not isinstance(origin, DecoratableObject):
+        mod = inspect.getmodule(origin)
+
+        name = origin.__name__
+        origin = DecoratableObject(origin)
+        setattr(mod, name, origin)
+
+    return _create_decorator(origin, needs_origin, mode)
+
+
+def decorate(origin, needs_origin=True, mode=DecoratableObject.MODE_OVERRIDE):
+    if inspect.ismethod(origin):
+        return _decorate_method(origin, needs_origin, mode)
+
+    if inspect.isfunction(origin) or isinstance(origin, DecoratableObject):
+        return _decorate_function(origin, needs_origin, mode)
+
+    def decorator(fn):
+        return fn
 
 
+def _decorate_params(origin):
+    return decorate(origin, mode=DecoratableObject.MODE_PARAMS)
 
+decorate.params = _decorate_params
 
+def _decorate_result(origin):
+    return decorate(origin, mode=DecoratableObject.MODE_RESULT)
 
+decorate.result = _decorate_result