Fix typing for airflow provider

yandex-cloud · Jul 22, 2024 · 9e93ea8 · 9e93ea8
1 parent 9bbc4da
commit 9e93ea8
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/yandexcloud/_operation_waiter.py b/yandexcloud/_operation_waiter.py
@@ -56,11 +56,11 @@ def get_operation_result(
  meta_type: Optional[Type["MetaType"]] = None,
  timeout: Optional[float] = None,
  logger: Optional[logging.Logger] = None,
-) -> Union["OperationResult", "OperationError"]:
+) -> Union["OperationResult[ResponseType, MetaType]", "OperationError"]:
  if not logger:
  logger = logging.getLogger()
  logger.addHandler(logging.NullHandler())
- operation_result = OperationResult(operation)
+ operation_result: OperationResult[ResponseType, MetaType] = OperationResult(operation)
  created_at = datetime.fromtimestamp(operation.created_at.seconds)
  message = (
  "Running Yandex.Cloud operation. ID: {id}. "

diff --git a/yandexcloud/_sdk.py b/yandexcloud/_sdk.py
@@ -104,7 +104,7 @@ def wait_operation_and_get_result(
  meta_type: Optional[Type["MetaType"]] = None,
  timeout: Optional[float] = None,
  logger: Optional["logging.Logger"] = None,
- ) -> Union["OperationResult", "OperationError"]:
+ ) -> Union["OperationResult[ResponseType, MetaType]", "OperationError"]:
  return _operation_waiter.get_operation_result(self, operation, response_type, meta_type, timeout, logger)
 
  def create_operation_and_get_result(

diff --git a/yandexcloud/operations.py b/yandexcloud/operations.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, TypeVar
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar
 
 if TYPE_CHECKING:
  import google.protobuf.message
@@ -11,7 +11,7 @@
 MetaType = TypeVar("MetaType", bound="google.protobuf.message.Message") # pylint: disable=C0103
 
 
-class OperationResult:
+class OperationResult(Generic[ResponseType, MetaType]):
  def __init__(
  self,
  operation: "Operation",
@@ -24,7 +24,7 @@ def __init__(
 
 
 class OperationError(RuntimeError):
- def __init__(self, message: str, operation_result: OperationResult):
+ def __init__(self, message: str, operation_result: OperationResult[ResponseType, MetaType]):
  super(OperationError, self).__init__(message) # pylint: disable=super-with-arguments
  self.message = message
  self.operation_result = operation_result