From 781397163651e26297a8c24ec9460451adeab490 Mon Sep 17 00:00:00 2001
From: Zhiyuan Chen <this@zyc.ai>
Date: Tue, 25 Apr 2023 16:57:41 +0800
Subject: [PATCH] remove lambda in registry.register to support serialisation

Signed-off-by: Zhiyuan Chen <this@zyc.ai>
---
 chanfig/registry.py | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/chanfig/registry.py b/chanfig/registry.py
index ec8326c2..21ff671c 100644
--- a/chanfig/registry.py
+++ b/chanfig/registry.py
@@ -69,7 +69,7 @@ def __init__(self, override: bool = False):
         super().__init__()
         self.setattr("override", override)
 
-    def register(self, component: Optional[Callable] = None, name: Optional[str] = None) -> Callable:
+    def register(self, component: Any = None, name: Optional[Any] = None) -> Callable:
         r"""
         Register a new component.
 
@@ -79,6 +79,7 @@ def register(self, component: Optional[Callable] = None, name: Optional[str] = N
 
         Returns:
             component: The registered component.
+                Registered component are expected to be `Callable`.
 
         Raises:
             ValueError: If the component with the same name already registered and `Registry.override=False`.
@@ -109,20 +110,25 @@ def register(self, component: Optional[Callable] = None, name: Optional[str] = N
         # Registry.register()
         if name is not None:
             self.set(name, component)
-
-        # @Registry.register()
-        @wraps(self.register)
-        def register(component, name=None):
-            if name is None:
-                name = component.__name__
-            self.set(name, component)
-            return component
-
+            return component  # type: ignore
         # @Registry.register
-        if callable(component) and name is None:
-            return register(component)
+        elif callable(component) and name is None:
+            self.set(component.__name__, component)
+            return component
 
-        return lambda x: register(x, component)
+        # @Registry.register()
+        def decorator(name: Optional[Any] = None):
+            @wraps(self.register)
+            def wrapper(component):
+                if name is None:
+                    self.set(component.__name__, component)
+                else:
+                    self.set(name, component)
+                return component
+
+            return wrapper
+
+        return decorator(component)
 
     def lookup(self, name: str) -> Any:
         r"""