1
0
mirror of https://github.com/dkmstr/openuds.git synced 2025-01-03 01:17:56 +03:00

Refactor ensure_connected decorator to support classes with connect method

This commit is contained in:
Adolfo Gómez García 2024-10-12 07:04:32 +02:00
parent c041284588
commit 203a46a804
No known key found for this signature in database
GPG Key ID: DD1ABF20724CDA23

View File

@ -138,24 +138,28 @@ def deprecated_class_value(new_var_name: str) -> collections.abc.Callable[..., t
return functools.partial(innerDeprecated, newVarName=new_var_name) return functools.partial(innerDeprecated, newVarName=new_var_name)
# Keep this, but mypy does not likes it... it's perfect with pyright
# # So only classes that have a "connect" method can use this decorator # # So only classes that have a "connect" method can use this decorator
class _HasConnect(typing.Protocol): class _HasConnect(typing.Protocol):
def connect(self) -> None: ... def connect(self) -> None: ...
# HasConnect = typing.TypeVar('HasConnect', bound=_HasConnect) # def ensure_connected(func: collections.abc.Callable[P, T]) -> collections.abc.Callable[P, T]:
# def ensure_connected(func: collections.abc.Callable[typing.Concatenate[HasConnect, P], T]) -> collections.abc.Callable[typing.Concatenate[HasConnect, P], T]:
# Keep this, but mypy does not likes it... it's perfect with pyright
# We use pyright for type checking, so we will use this
HasConnect = typing.TypeVar('HasConnect', bound=_HasConnect)
def ensure_connected(func: collections.abc.Callable[P, T]) -> collections.abc.Callable[P, T]: def ensure_connected(
func: collections.abc.Callable[typing.Concatenate[HasConnect, P], T]
) -> collections.abc.Callable[typing.Concatenate[HasConnect, P], T]:
"""This decorator calls "connect" method of the class of the wrapped object""" """This decorator calls "connect" method of the class of the wrapped object"""
@functools.wraps(func) @functools.wraps(func)
def new_func(*args: P.args, **kwargs: P.kwargs) -> T: def new_func(obj: HasConnect, /, *args: P.args, **kwargs: P.kwargs) -> T:
self = typing.cast(_HasConnect, args[0]) # self = typing.cast(_HasConnect, args[0])
self.connect() obj.connect()
return func(*args, **kwargs) return func(obj, *args, **kwargs)
return new_func return new_func