fix model directory typing

Signed-off-by: weiwee <wbwmat@gmail.com>
FederatedAI · Jun 28, 2023 · 343741f · 343741f
1 parent bf049d8
commit 343741f
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/python/fate/components/core/component_desc/artifacts/model/__init__.py b/python/fate/components/core/component_desc/artifacts/model/__init__.py
@@ -1,7 +1,11 @@
 from typing import Iterator, List, Optional, Type
 
 from .._base_type import Role, _create_artifact_annotation
-from ._directory import ModelDirectoryArtifactDescribe
+from ._directory import (
+    ModelDirectoryArtifactDescribe,
+    ModelDirectoryReader,
+    ModelDirectoryWriter,
+)
 from ._json import JsonModelArtifactDescribe, JsonModelReader, JsonModelWriter
 
 
@@ -21,25 +25,21 @@ def json_model_outputs(roles: Optional[List[Role]] = None, desc="", optional=Fal
     return _create_artifact_annotation(False, True, JsonModelArtifactDescribe, "model")(roles, desc, optional)
 
 
-def model_directory_input(
-    roles: Optional[List[Role]] = None, desc="", optional=False
-) -> Type[ModelDirectoryArtifactDescribe]:
+def model_directory_input(roles: Optional[List[Role]] = None, desc="", optional=False) -> Type[ModelDirectoryReader]:
     return _create_artifact_annotation(True, False, ModelDirectoryArtifactDescribe, "model")(roles, desc, optional)
 
 
 def model_directory_inputs(
     roles: Optional[List[Role]] = None, desc="", optional=False
-) -> Type[List[ModelDirectoryArtifactDescribe]]:
+) -> Type[List[ModelDirectoryReader]]:
     return _create_artifact_annotation(True, True, ModelDirectoryArtifactDescribe, "model")(roles, desc, optional)
 
 
-def model_directory_output(
-    roles: Optional[List[Role]] = None, desc="", optional=False
-) -> Type[ModelDirectoryArtifactDescribe]:
+def model_directory_output(roles: Optional[List[Role]] = None, desc="", optional=False) -> Type[ModelDirectoryWriter]:
     return _create_artifact_annotation(False, False, ModelDirectoryArtifactDescribe, "model")(roles, desc, optional)
 
 
 def model_directory_outputs(
     roles: Optional[List[Role]] = None, desc="", optional=False
-) -> Type[Iterator[ModelDirectoryArtifactDescribe]]:
+) -> Type[Iterator[ModelDirectoryWriter]]:
     return _create_artifact_annotation(False, True, ModelDirectoryArtifactDescribe, "model")(roles, desc, optional)