From 6d897f220cd509b6fdfeffa1f49190651bd797b1 Mon Sep 17 00:00:00 2001
From: deptyped <deptyped@gmail.com>
Date: Wed, 8 Jul 2020 19:30:35 +0300
Subject: [PATCH] Add ability to use custom method name

---
 jsonrpcserver/methods.py | 19 +++++++++++++++++++
 tests/test_methods.py    | 10 ++++++++++
 2 files changed, 29 insertions(+)

diff --git a/jsonrpcserver/methods.py b/jsonrpcserver/methods.py
index 67fe6c2..e830368 100644
--- a/jsonrpcserver/methods.py
+++ b/jsonrpcserver/methods.py
@@ -65,7 +65,26 @@ def add(self, *args: Any, **kwargs: Any) -> Optional[Callable]:
             @methods.add
             def subtract(minuend, subtrahend):
                 return minuend - subtrahend
+
+            @methods.add(name='divide')
+            def division(dividend, divisor):
+                return dividend / divisor
         """
+        if "name" in kwargs and isinstance(kwargs["name"], str):
+            return self._parameterized_add(*args, **kwargs)
+        else:
+            return self._batch_add(*args, **kwargs)
+
+    def _parameterized_add(self, name: str) -> Callable:
+        def decorator(method):
+            assert callable(method)
+            self.items[name] = method
+
+            return method
+
+        return decorator
+
+    def _batch_add(self, *args: Any, **kwargs: Any) -> Optional[Callable]:
         # Multiple loops here, but due to changes in dictionary comprehension evaluation
         # order in Python 3.8 (PEP 572), we need to validate separately from the
         # dictionary comprehension. Otherwise different exceptions will be raised in 3.8
diff --git a/tests/test_methods.py b/tests/test_methods.py
index 92786c5..6d3dc62 100644
--- a/tests/test_methods.py
+++ b/tests/test_methods.py
@@ -148,6 +148,16 @@ def foo():
     assert methods.items["foo"] is foo
 
 
+def test_add_function_custom_name_via_decorator():
+    methods = Methods()
+
+    @methods.add(name='bar')
+    def foo():
+        pass
+
+    assert methods.items["bar"] is foo
+
+
 def test_add_static_method_via_decorator():
     methods = Methods()