fix bug

Signed-off-by: sagewe <wbwmat@gmail.com>
FederatedAI · Dec 11, 2023 · 98a72f0 · 98a72f0
1 parent bf9e8a6
commit 98a72f0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/python/fate/arch/computing/api/_profile.py b/python/fate/arch/computing/api/_profile.py
@@ -339,7 +339,10 @@ def _call_stack_strings():
     return call_stack_strings
 
 
-def computing_profile(func):
+T = typing.TypeVar("T", bound=typing.Callable)
+
+
+def computing_profile(func: T) -> T:
     @wraps(func)
     def _fn(*args, **kwargs):
         function_call_stack = _call_stack_strings()

diff --git a/python/fate/arch/context/_context.py b/python/fate/arch/context/_context.py
@@ -111,7 +111,7 @@ def with_namespace(self, namespace: NS):
         )
 
     @property
-    def computing(self):
+    def computing(self) -> "KVTableContext":
         return self._get_computing()
 
     @property
@@ -259,7 +259,7 @@ def _get_federation(self):
             raise RuntimeError(f"federation not set")
         return self._federation
 
-    def _get_computing(self):
+    def _get_computing(self) -> "KVTableContext":
         if self._computing is None:
             raise RuntimeError(f"computing not set")
         return self._computing