From 177f59ea6c7ca6f3d9d433826b946bc96462cc42 Mon Sep 17 00:00:00 2001
From: 苏向夜 <fu050409@163.com>
Date: Fri, 15 Mar 2024 14:54:00 +0800
Subject: feat(register): add namespace field in register

---
 src/infini/register.py | 27 ++++++++++++++++++++++++---
 src/infini/typing.py   |  1 +
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/infini/register.py b/src/infini/register.py
index 4b191cb6..8ccff3d6 100644
--- a/src/infini/register.py
+++ b/src/infini/register.py
@@ -22,7 +22,13 @@ class Register:
         self.interceptors = []
         self.generators = {}
 
-    def pre_interceptor(self, router: Union[Router, str], priority: int = 0):
+    def pre_interceptor(
+        self,
+        router: Union[Router, str],
+        *,
+        priority: int = 0,
+        namespace: Optional[str] = None,
+    ):
         def decorator(func):
             @wraps(func)
             def wrapper(*args, **kwargs) -> Union[Input, Output]:
@@ -33,13 +39,20 @@ class Register:
                     "priority": priority,
                     "router": Contains(router) if isinstance(router, str) else router,
                     "handler": wrapper,
+                    "namespace": namespace,
                 }
             )
             return wrapper
 
         return decorator
 
-    def handler(self, router: Union[Router, str], priority: int = 0):
+    def handler(
+        self,
+        router: Union[Router, str],
+        *,
+        priority: int = 0,
+        namespace: Optional[str] = None,
+    ):
         def decorator(func):
             @wraps(func)
             def wrapper(*args, **kwargs) -> Output:
@@ -50,6 +63,7 @@ class Register:
                     "priority": priority,
                     "router": Contains(router) if isinstance(router, str) else router,
                     "handler": wrapper,
+                    "namespace": namespace,
                 }
             )
             return wrapper
@@ -73,7 +87,13 @@ class Register:
 
         return decorator
 
-    def interceptor(self, router: Union[Router, str], priority: int = 0):
+    def interceptor(
+        self,
+        router: Union[Router, str],
+        *,
+        priority: int = 0,
+        namespace: Optional[str] = None,
+    ):
         def decorator(func):
             @wraps(func)
             def wrapper(*args, **kwargs) -> Union[Input, Output]:
@@ -84,6 +104,7 @@ class Register:
                     "priority": priority,
                     "router": Contains(router) if isinstance(router, str) else router,
                     "handler": wrapper,
+                    "namespace": namespace,
                 }
             )
             return wrapper
diff --git a/src/infini/typing.py b/src/infini/typing.py
index 14c219fc..24faad83 100644
--- a/src/infini/typing.py
+++ b/src/infini/typing.py
@@ -31,3 +31,4 @@ class RouterType(TypedDict):
     priority: int
     router: "Router"
     handler: Callable[..., Union[Stream, OutputGenerator]]
+    namespace: Optional[str]
-- 
cgit v1.2.3-70-g09d2