Proper typing of the decorator

src/zenml/client.py

@@ -188,6 +188,7 @@
 logger = get_logger(__name__)
 
 AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
+T = TypeVar("T")
 
 
 class ClientConfiguration(FileSyncModel):
@@ -365,7 +366,7 @@ def _reset_instance(cls, client: Optional["Client"] = None) -> None:
         cls._global_client = client
 
     @staticmethod
-    def _fail_for_sql_zen_store(method):
+    def _fail_for_sql_zen_store(method: Callable[..., T]) -> Callable[..., T]:
         """Decorator for all methods, that are disallowed when the client is not connected through REST API.
 
         Args:
@@ -376,7 +377,7 @@ def _fail_for_sql_zen_store(method):
         """
 
         @functools.wraps(method)
-        def wrapper(self, *args, **kwargs):
+        def wrapper(self: "Client", *args: Any, **kwargs: Any) -> Any:
             # No isinstance check to avoid importing ZenStore implementations
             if self.zen_store.__class__.__name__ == "SqlZenStore":
                 raise TypeError(