update

Python-Tools · Jun 18, 2024 · 7f17c93 · 7f17c93
1 parent 84e5955
commit 7f17c93
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/schema_entry/entrypoint.py b/schema_entry/entrypoint.py
@@ -158,14 +158,14 @@ def regist_sub(self, subcmdclz: type, **kwargs: Any) -> EntryPointABC:
         return instance
 
     def as_main(self, func: Callable[..., Optional[Any]]) -> Callable[..., Optional[Any]]:
-        @ functools.wraps(func)
+        @functools.wraps(func)
         def warp(*args: Any, **kwargs: Any) -> Optional[Any]:
             return func(*args, **kwargs)
 
         self._main = warp
         return warp
 
-    def with_schema(self, schemaObj: Union[str, dict, PydanticModelLike]) -> None:
+    def with_schema(self, schemaObj: Union[str, dict, PydanticModelLike]) -> Union[str, dict, PydanticModelLike]:
         """注册schema
 
         # todo 注册schema或pydantic类
@@ -182,6 +182,7 @@ def with_schema(self, schemaObj: Union[str, dict, PydanticModelLike]) -> None:
             self.schema = pydantic_schema_to_protocol(schema)
             self._name = schemaObj.__name__
             self.__doc__ = schemaObj.__doc__
+        return schemaObj
 
     def __call__(self, argv: Sequence[str]) -> Optional[Any]:
         if not self.usage:

diff --git a/schema_entry/entrypoint_base.py b/schema_entry/entrypoint_base.py
@@ -143,7 +143,7 @@ def as_main(self, func: Callable[..., Optional[Any]]) -> Callable[..., Optional[
         """
 
     @abc.abstractmethod
-    def with_schema(self, schemaObj: Union[str, dict, PydanticModelLike]) -> None:
+    def with_schema(self, schemaObj: Union[str, dict, PydanticModelLike]) -> Union[str, dict, PydanticModelLike]:
         """注册schema
 
         可以是一个json字符串,一个dict,或一个pydantic的类