From 0d29820555b6cca701ac2fd43e3b49121e44fb7e Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Fri, 31 May 2024 17:56:53 +0300 Subject: [PATCH] fix: Avoid loosing type hints because of decorators (#379) Co-authored-by: Diwank Singh Tomer --- sdks/python/julep/managers/utils.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/sdks/python/julep/managers/utils.py b/sdks/python/julep/managers/utils.py index 58d198671..5123f102c 100644 --- a/sdks/python/julep/managers/utils.py +++ b/sdks/python/julep/managers/utils.py @@ -2,9 +2,14 @@ from functools import wraps from typing import Callable from uuid import UUID +from typing_extensions import ParamSpec +from pydantic.main import Model from ..api.types import ResourceCreatedResponse +P = ParamSpec("P") + + class NotSet: pass @@ -33,19 +38,20 @@ def is_valid_uuid4(uuid_to_test: str) -> bool: return True -def rewrap_in_class(cls): - def decorator(func: Callable[..., ResourceCreatedResponse]): - @wraps(func) +def rewrap_in_class(cls: type[Model]): + def decorator(func: Callable[P, ResourceCreatedResponse]): # This wrapper is used for asynchronous functions to ensure they are properly awaited and their results are processed by `cls.construct`. + @wraps(func) async def async_wrapper(*args, **kwargs): result = await func(*args, **kwargs) - return cls.construct(**kwargs, **result.dict()) + return cls.model_construct(**kwargs, **result.dict()) # This wrapper handles synchronous functions, directly calling them and processing their results with `cls.construct`. + @wraps(func) def sync_wrapper(*args, **kwargs): # Logging at this point might be useful for debugging, but should use a proper logging framework instead of print statements for production code. result = func(*args, **kwargs) - return cls.construct(**kwargs, **result.dict()) + return cls.model_construct(**kwargs, **result.dict()) return async_wrapper if iscoroutinefunction(func) else sync_wrapper