diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py index 174af1e39e2cc..a063b39b1a35f 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py @@ -34,11 +34,10 @@ from typing import IO, TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload from urllib.parse import urlsplit +# Make mypy happy by importing as aliases +import google.cloud.storage as storage from gcloud.aio.storage import Storage from google.api_core.exceptions import GoogleAPICallError, NotFound - -# not sure why but mypy complains on missing `storage` but it is clearly there and is importable -from google.cloud import storage # type: ignore[attr-defined] from google.cloud.exceptions import GoogleCloudError from google.cloud.storage.retry import DEFAULT_RETRY diff --git a/providers/google/src/airflow/providers/google/cloud/log/gcs_task_handler.py b/providers/google/src/airflow/providers/google/cloud/log/gcs_task_handler.py index 894f29138d0a0..6c30f78a6a828 100644 --- a/providers/google/src/airflow/providers/google/cloud/log/gcs_task_handler.py +++ b/providers/google/src/airflow/providers/google/cloud/log/gcs_task_handler.py @@ -27,8 +27,8 @@ import attrs -# not sure why but mypy complains on missing `storage` but it is clearly there and is importable -from google.cloud import storage # type: ignore[attr-defined] +# Make mypy happy by importing as aliases +import google.cloud.storage as storage from airflow.configuration import conf from airflow.exceptions import AirflowNotFoundException diff --git a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py index c4ba58fdb24b8..5126a34ba6e37 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py @@ -28,11 +28,10 @@ from unittest.mock import MagicMock import dateutil +import google.cloud.storage as storage import pytest from google.api_core.exceptions import GoogleAPICallError - -# dynamic storage type in google.cloud needs to be type-ignored -from google.cloud import exceptions, storage # type: ignore[attr-defined] +from google.cloud.exceptions import NotFound from google.cloud.storage.retry import DEFAULT_RETRY from airflow.exceptions import AirflowException @@ -559,9 +558,9 @@ def test_delete_nonexisting_object(self, mock_service): bucket_method = mock_service.return_value.bucket blob = bucket_method.return_value.blob delete_method = blob.return_value.delete - delete_method.side_effect = exceptions.NotFound(message="Not Found") + delete_method.side_effect = NotFound(message="Not Found") - with pytest.raises(exceptions.NotFound): + with pytest.raises(NotFound): self.gcs_hook.delete(bucket_name=test_bucket, object_name=test_object) @mock.patch(GCS_STRING.format("GCSHook.get_conn")) @@ -598,9 +597,7 @@ def test_delete_bucket(self, mock_service): @mock.patch(GCS_STRING.format("GCSHook.get_conn")) def test_delete_nonexisting_bucket(self, mock_service, caplog): - mock_service.return_value.bucket.return_value.delete.side_effect = exceptions.NotFound( - message="Not Found" - ) + mock_service.return_value.bucket.return_value.delete.side_effect = NotFound(message="Not Found") test_bucket = "test bucket" with caplog.at_level(logging.INFO): self.gcs_hook.delete_bucket(bucket_name=test_bucket)