[refactor] make node a non-opional attr for AuthedServiceContext

reenable checking for mypy error `union-attr`
OpenMined · Feb 26, 2024 · c827525 · c827525
1 parent beab7ca
commit c827525
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 5 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -190,7 +190,6 @@ repos:
             "--install-types",
             "--non-interactive",
             "--config-file=tox.ini",
-            "--disable-error-code=union-attr", # todo: remove this line after fixing the issue context.node can be None
           ]
 
   - repo: https://github.com/kynan/nbstripout

diff --git a/packages/syft/src/syft/service/code/user_code.py b/packages/syft/src/syft/service/code/user_code.py
@@ -944,14 +944,15 @@ def syft_function(
         output_policy_type = type(output_policy)
 
     def decorator(f: Any) -> SubmitUserCode:
+        # TODO: fix the mypy issue below
         res = SubmitUserCode(
             code=inspect.getsource(f),
             func_name=f.__name__,
             signature=inspect.signature(f),
             input_policy_type=input_policy_type,
-            input_policy_init_kwargs=input_policy.init_kwargs,
+            input_policy_init_kwargs=input_policy.init_kwargs,  # type: ignore
             output_policy_type=output_policy_type,
-            output_policy_init_kwargs=output_policy.init_kwargs,
+            output_policy_init_kwargs=output_policy.init_kwargs,  # type: ignore
             local_function=f,
             input_kwargs=f.__code__.co_varnames[: f.__code__.co_argcount],
             worker_pool_name=worker_pool_name,

diff --git a/packages/syft/src/syft/service/context.py b/packages/syft/src/syft/service/context.py
@@ -36,6 +36,7 @@ class AuthedServiceContext(NodeServiceContext):
     job_id: Optional[UID]
     extra_kwargs: Dict = {}
     has_execute_permissions: bool = False
+    node: AbstractNode
 
     def capabilities(self) -> List[ServiceRoleCapability]:
         return ROLE_TO_CAPABILITIES.get(self.role, [])

diff --git a/packages/syft/src/syft/service/network/node_peer.py b/packages/syft/src/syft/service/network/node_peer.py
@@ -2,6 +2,7 @@
 from typing import List
 from typing import Optional
 from typing import Tuple
+from typing import Union
 
 # third party
 from typing_extensions import Self
@@ -115,7 +116,9 @@ def from_client(cls, client: SyftClient) -> Self:
         peer.node_routes.append(route)
         return peer
 
-    def client_with_context(self, context: NodeServiceContext) -> SyftClient:
+    def client_with_context(
+        self, context: NodeServiceContext
+    ) -> Union[SyftClient, SyftError]:
         if len(self.node_routes) < 1:
             raise Exception(f"No routes to peer: {self}")
         # select the latest added route
@@ -125,6 +128,8 @@ def client_with_context(self, context: NodeServiceContext) -> SyftClient:
         client_type = connection.get_client_type()
         if isinstance(client_type, SyftError):
             return client_type
+        if context.node is None:
+            return SyftError(message=f"context {context}'s node is None")
         return client_type(connection=connection, credentials=context.node.signing_key)
 
     def client_with_key(self, credentials: SyftSigningKey) -> SyftClient:

diff --git a/packages/syft/src/syft/service/network/routes.py b/packages/syft/src/syft/service/network/routes.py
@@ -33,11 +33,15 @@
 
 
 class NodeRoute:
-    def client_with_context(self, context: NodeServiceContext) -> SyftClient:
+    def client_with_context(
+        self, context: NodeServiceContext
+    ) -> Union[SyftClient, SyftError]:
         connection = route_to_connection(route=self, context=context)
         client_type = connection.get_client_type()
         if isinstance(client_type, SyftError):
             return client_type
+        if context.node is None:
+            return SyftError(message=f"context {context}'s node is None")
         return client_type(connection=connection, credentials=context.node.signing_key)
 
     def validate_with_context(self, context: AuthedServiceContext) -> NodePeer:

diff --git a/packages/syft/src/syft/service/service.py b/packages/syft/src/syft/service/service.py
@@ -70,6 +70,8 @@ def resolve_link(
         if isinstance(obj, OkErr) and obj.is_ok():
             obj = obj.ok()
         if hasattr(obj, "node_uid"):
+            if context.node is None:
+                return SyftError(message=f"context {context}'s node is None")
             obj.node_uid = context.node.id
         if not isinstance(obj, OkErr):
             obj = Ok(obj)