Skip to content

Commit 68a6a05

Browse files
authored
Fix mypy checks for new azure libraries (#41386)
1 parent 54c165c commit 68a6a05

File tree

1 file changed

+9
-5
lines changed
  • airflow/providers/microsoft/azure/hooks

1 file changed

+9
-5
lines changed

airflow/providers/microsoft/azure/hooks/wasb.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,8 @@ def get_conn(self) -> BlobServiceClient:
213213
**extra,
214214
)
215215

216-
def _get_container_client(self, container_name: str) -> ContainerClient:
216+
# TODO: rework the interface as it might also return AsyncContainerClient
217+
def _get_container_client(self, container_name: str) -> ContainerClient: # type: ignore[override]
217218
"""
218219
Instantiate a container client.
219220
@@ -222,7 +223,7 @@ def _get_container_client(self, container_name: str) -> ContainerClient:
222223
"""
223224
return self.blob_service_client.get_container_client(container_name)
224225

225-
def _get_blob_client(self, container_name: str, blob_name: str) -> BlobClient:
226+
def _get_blob_client(self, container_name: str, blob_name: str) -> BlobClient | AsyncBlobClient:
226227
"""
227228
Instantiate a blob client.
228229
@@ -415,7 +416,8 @@ def upload(
415416
self.create_container(container_name)
416417

417418
blob_client = self._get_blob_client(container_name, blob_name)
418-
return blob_client.upload_blob(data, blob_type, length=length, **kwargs)
419+
# TODO: rework the interface as it might also return Awaitable
420+
return blob_client.upload_blob(data, blob_type, length=length, **kwargs) # type: ignore[return-value]
419421

420422
def download(
421423
self, container_name, blob_name, offset: int | None = None, length: int | None = None, **kwargs
@@ -430,7 +432,8 @@ def download(
430432
:param length: Number of bytes to read from the stream.
431433
"""
432434
blob_client = self._get_blob_client(container_name, blob_name)
433-
return blob_client.download_blob(offset=offset, length=length, **kwargs)
435+
# TODO: rework the interface as it might also return Awaitable
436+
return blob_client.download_blob(offset=offset, length=length, **kwargs) # type: ignore[return-value]
434437

435438
def create_container(self, container_name: str) -> None:
436439
"""
@@ -656,7 +659,8 @@ async def check_for_blob_async(self, container_name: str, blob_name: str, **kwar
656659
return False
657660
return True
658661

659-
def _get_container_client(self, container_name: str) -> AsyncContainerClient:
662+
# TODO: rework the interface as in parent Hook it returns ContainerClient
663+
def _get_container_client(self, container_name: str) -> AsyncContainerClient: # type: ignore[override]
660664
"""
661665
Instantiate a container client.
662666

0 commit comments

Comments
 (0)