diff --git a/jsonrpc/backend/flask.py b/jsonrpc/backend/flask.py index d73c273..943f175 100644 --- a/jsonrpc/backend/flask.py +++ b/jsonrpc/backend/flask.py @@ -19,7 +19,8 @@ class JSONRPCAPI(object): - def __init__(self, dispatcher=None, check_content_type=True): + def __init__(self, dispatcher=None, check_content_type=True, + json_encoder=None): """ :param dispatcher: methods dispatcher @@ -31,6 +32,7 @@ def __init__(self, dispatcher=None, check_content_type=True): self.dispatcher = dispatcher if dispatcher is not None \ else Dispatcher() self.check_content_type = check_content_type + self.json_encoder = json_encoder or DatetimeDecimalEncoder def as_blueprint(self, name=None): blueprint = Blueprint(name if name else str(uuid4()), __name__) @@ -83,9 +85,9 @@ def _get_request_str(self): return request.data return list(request.form.keys())[0] - @staticmethod - def _serialize(s): - return json.dumps(s, cls=DatetimeDecimalEncoder) + # @staticmethod + def _serialize(self, s): + return json.dumps(s, cls=self.json_encoder) api = JSONRPCAPI() diff --git a/jsonrpc/dispatcher.py b/jsonrpc/dispatcher.py index 2bf5e43..2b42838 100644 --- a/jsonrpc/dispatcher.py +++ b/jsonrpc/dispatcher.py @@ -3,7 +3,9 @@ For usage examples see :meth:`Dispatcher.add_method` """ +from . import six import collections +from functools import wraps class Dispatcher(collections.MutableMapping): @@ -27,7 +29,10 @@ def __init__(self, prototype=None): None """ + self._decorators = [] self.method_map = dict() + self._before_request_hooks = [] + self._error_handler_spec = {} if prototype is not None: self.build_method_map(prototype) @@ -36,7 +41,7 @@ def __getitem__(self, key): return self.method_map[key] def __setitem__(self, key, value): - self.method_map[key] = value + self.method_map[key] = self._wrap_method(value) def __delitem__(self, key): del self.method_map[key] @@ -50,10 +55,47 @@ def __iter__(self): def __repr__(self): return repr(self.method_map) + def register_decorator(self, a): + self._decorators.extend(a if hasattr(a, '__iter__') else [a]) + + def before_request(self, hook): + self._before_request_hooks.append(hook) + + def errorhandler(self, exception): + def decorator(f): + self._error_handler_spec[exception] = f + return f + return decorator + + def _wrap_method(self, f): + @wraps(f) + def _method(*args, **kwargs): + try: + for hook in self._before_request_hooks: + hook() + + nf = f + for deco in reversed(self._decorators): + nf = deco(nf) + + return nf(*args, **kwargs) + except Exception as e: + for E, h in self._error_handler_spec.items(): + if isinstance(e, E): + return h(e) + raise + + return _method + def add_class(self, cls): - prefix = cls.__name__.lower() + '.' + if hasattr(cls, 'rpc_method_prefix'): + prefix = cls.rpc_method_prefix + '.' + else: + prefix = cls.__name__.lower() + '.' self.build_method_map(cls(), prefix) + return cls # for working as decorator + def add_object(self, obj): prefix = obj.__class__.__name__.lower() + '.' self.build_method_map(obj, prefix) @@ -93,8 +135,21 @@ def add_method(self, f, name=None): def mymethod(*args, **kwargs): print(args, kwargs) + >>> @d.add_method(name='MyMethod') + def mymethod(*args, **kwargs): + print(args, kwargs) + """ - self.method_map[name or f.__name__] = f + if isinstance(f, six.string_types): + name, f = f, name + + if f is None: + # Be decorator generator + def _add_method(f): + return self.add_method(f, name) + return _add_method + + self[name or f.__name__] = f return f def build_method_map(self, prototype, prefix=''): @@ -112,10 +167,19 @@ def build_method_map(self, prototype, prefix=''): Prefix of methods """ + rpc_exports = getattr(prototype, 'rpc_exports', None) + + def _should_export(method): + if method.startswith('_'): + return False + if rpc_exports is None: + return True + return method in rpc_exports + if not isinstance(prototype, dict): prototype = dict((method, getattr(prototype, method)) for method in dir(prototype) - if not method.startswith('_')) + if _should_export(method)) for attr, method in prototype.items(): if callable(method): diff --git a/jsonrpc/tests/test_dispatcher.py b/jsonrpc/tests/test_dispatcher.py index 56876b2..b74f414 100644 --- a/jsonrpc/tests/test_dispatcher.py +++ b/jsonrpc/tests/test_dispatcher.py @@ -90,10 +90,10 @@ def test_to_dict(self): d = Dispatcher() def func(): - return "" + return "x" d["method"] = func - self.assertEqual(dict(d), {"method": func}) + self.assertEqual(dict(d)["method"](), "x") def test_init_from_object_instance(self):