[Storage] Adjust return type of `StorageStreamDownloader.readall` (#25174)

sdk/storage/azure-storage-blob/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Features Added
 
 ### Bugs Fixed
+- Adjusted type hints for `upload_blob` and `StorageStreamDownloader.readall`.
 
 ## 12.13.0 (2022-07-07)
 

sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py

@@ -8,7 +8,7 @@
 from functools import partial
 from io import BytesIO
 from typing import (
-    Any, AnyStr, Dict, IO, Iterable, List, Optional, Tuple, Type, TypeVar, Union,
+    Any, AnyStr, Dict, IO, Iterable, List, Optional, overload, Tuple, Type, TypeVar, Union,
     TYPE_CHECKING
 )
 from urllib.parse import urlparse, quote, unquote
@@ -733,8 +733,8 @@ class BlobClient(StorageAccountHostsMixin, StorageEncryptionMixin):  # pylint: d
             return upload_page_blob(**options)
         return upload_append_blob(**options)
 
-    def _download_blob_options(self, offset=None, length=None, **kwargs):
-        # type: (Optional[int], Optional[int], **Any) -> Dict[str, Any]
+    def _download_blob_options(self, offset=None, length=None, encoding=None, **kwargs):
+        # type: (Optional[int], Optional[int], Optional[str], **Any) -> Dict[str, Any]
         if self.require_encryption and not self.key_encryption_key:
             raise ValueError("Encryption required but no key was provided.")
         if length is not None and offset is None:
@@ -768,18 +768,40 @@ class BlobClient(StorageAccountHostsMixin, StorageEncryptionMixin):  # pylint: d
             'lease_access_conditions': access_conditions,
             'modified_access_conditions': mod_conditions,
             'cpk_info': cpk_info,
-            'cls': kwargs.pop('cls', None) or deserialize_blob_stream,
+            'download_cls': kwargs.pop('cls', None) or deserialize_blob_stream,
             'max_concurrency':kwargs.pop('max_concurrency', 1),
-            'encoding': kwargs.pop('encoding', None),
+            'encoding': encoding,
             'timeout': kwargs.pop('timeout', None),
             'name': self.blob_name,
             'container': self.container_name}
         options.update(kwargs)
         return options
 
+    @overload
+    def download_blob(
+            self, offset: int = None,
+            length: int = None,
+            *,
+            encoding: str,
+            **kwargs) -> StorageStreamDownloader[str]:
+        ...
+
+    @overload
+    def download_blob(
+            self, offset: int = None,
+            length: int = None,
+            *,
+            encoding: None = None,
+            **kwargs) -> StorageStreamDownloader[bytes]:
+        ...
+
     @distributed_trace
-    def download_blob(self, offset=None, length=None, **kwargs):
-        # type: (Optional[int], Optional[int], **Any) -> StorageStreamDownloader
+    def download_blob(
+            self, offset: int = None,
+            length: int = None,
+            *,
+            encoding: Optional[str] = None,
+            **kwargs) -> StorageStreamDownloader:
         """Downloads a blob to the StorageStreamDownloader. The readall() method must
         be used to read all the content or readinto() must be used to download the blob into
         a stream. Using chunks() returns an iterator which allows the user to iterate over the content in chunks.
@@ -867,6 +889,7 @@ class BlobClient(StorageAccountHostsMixin, StorageEncryptionMixin):  # pylint: d
         options = self._download_blob_options(
             offset=offset,
             length=length,
+            encoding=encoding,
             **kwargs)
         return StorageStreamDownloader(**options)
 

sdk/storage/azure-storage-blob/azure/storage/blob/_container_client.py

@@ -7,7 +7,7 @@
 
 import functools
 from typing import (  # pylint: disable=unused-import
-    Any, AnyStr, Dict, List, IO, Iterable, Iterator, Optional, TypeVar, Union,
+    Any, AnyStr, Dict, List, IO, Iterable, Iterator, Optional, overload, TypeVar, Union,
     TYPE_CHECKING
 )
 from urllib.parse import urlparse, quote, unquote
@@ -31,6 +31,7 @@ from ._generated import AzureBlobStorage
 from ._generated.models import SignedIdentifier
 from ._blob_client import BlobClient
 from ._deserialize import deserialize_container_properties
+from ._download import StorageStreamDownloader
 from ._encryption import StorageEncryptionMixin
 from ._lease import BlobLeaseClient
 from ._list_blobs_helper import BlobPrefix, BlobPropertiesPaged, FilteredBlobPaged
@@ -1060,9 +1061,34 @@ class ContainerClient(StorageAccountHostsMixin, StorageEncryptionMixin):    # py
             timeout=timeout,
             **kwargs)
 
+    @overload
+    def download_blob(
+            self, blob: Union[str, BlobProperties],
+            offset: int = None,
+            length: int = None,
+            *,
+            encoding: str,
+            **kwargs) -> StorageStreamDownloader[str]:
+        ...
+
+    @overload
+    def download_blob(
+            self, blob: Union[str, BlobProperties],
+            offset: int = None,
+            length: int = None,
+            *,
+            encoding: None = None,
+            **kwargs) -> StorageStreamDownloader[bytes]:
+        ...
+
     @distributed_trace
-    def download_blob(self, blob, offset=None, length=None, **kwargs):
-        # type: (Union[str, BlobProperties], Optional[int], Optional[int], **Any) -> StorageStreamDownloader
+    def download_blob(
+            self, blob: Union[str, BlobProperties],
+            offset: int = None,
+            length: int = None,
+            *,
+            encoding: Optional[str] = None,
+            **kwargs) -> StorageStreamDownloader:
         """Downloads a blob to the StorageStreamDownloader. The readall() method must
         be used to read all the content or readinto() must be used to download the blob into
         a stream. Using chunks() returns an iterator which allows the user to iterate over the content in chunks.
@@ -1143,7 +1169,11 @@ class ContainerClient(StorageAccountHostsMixin, StorageEncryptionMixin):    # py
         """
         blob_client = self.get_blob_client(blob) # type: ignore
         kwargs.setdefault('merge_span', True)
-        return blob_client.download_blob(offset=offset, length=length, **kwargs)
+        return blob_client.download_blob(
+            offset=offset,
+            length=length,
+            encoding=encoding,
+            **kwargs)
 
     def _generate_delete_blobs_subrequest_options(
         self, snapshot=None,

sdk/storage/azure-storage-blob/azure/storage/blob/_download.py

@@ -9,7 +9,7 @@ import threading
 import time
 import warnings
 from io import BytesIO
-from typing import Iterator, Union
+from typing import Generic, Iterator, TypeVar
 
 import requests
 from azure.core.exceptions import HttpResponseError, ServiceResponseError
@@ -26,6 +26,8 @@ from ._encryption import (
     parse_encryption_data
 )
 
+T = TypeVar('T', bytes, str)
+
 
 def process_range_and_offset(start_range, end_range, length, encryption_options, encryption_data):
     start_offset, end_offset = 0, 0
@@ -281,7 +283,7 @@ class _ChunkIterator(object):
         return chunk_data
 
 
-class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attributes
+class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-attributes
     """A streaming object to download from Azure Storage.
 
     :ivar str name:
@@ -308,6 +310,7 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
         name=None,
         container=None,
         encoding=None,
+        download_cls=None,
         **kwargs
     ):
         self.name = name
@@ -333,6 +336,10 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
         self._response = None
         self._encryption_data = None
 
+        # The cls is passed in via download_cls to avoid conflicting arg name with Generic.__new__
+        # but needs to be changed to cls in the request options.
+        self._request_options['cls'] = download_cls
+
         if self._encryption_options.get("key") is not None or self._encryption_options.get("resolver") is not None:
             self._get_encryption_data_request()
 
@@ -546,11 +553,11 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
             chunk_size=self._config.max_chunk_get_size)
 
     def readall(self):
-        # type: () -> Union[bytes, str]
+        # type: () -> T
         """Download the contents of this blob.
 
         This operation is blocking until all data is downloaded.
-        :rtype: bytes or str
+        :rtype: T
         """
         stream = BytesIO()
         self.readinto(stream)

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_blob_client_async.py

@@ -8,7 +8,7 @@
 import warnings
 from functools import partial
 from typing import (  # pylint: disable=unused-import
-    Any, AnyStr, Dict, IO, Iterable, List, Tuple, Optional, Union,
+    Any, AnyStr, Dict, IO, Iterable, List, Optional, overload, Tuple, Union,
     TYPE_CHECKING
 )
 
@@ -405,9 +405,31 @@ class BlobClient(AsyncStorageAccountHostsMixin, BlobClientBase, StorageEncryptio
             return await upload_page_blob(**options)
         return await upload_append_blob(**options)
 
+    @overload
+    async def download_blob(
+            self, offset: int = None,
+            length: int = None,
+            *,
+            encoding: str,
+            **kwargs) -> StorageStreamDownloader[str]:
+        ...
+
+    @overload
+    async def download_blob(
+            self, offset: int = None,
+            length: int = None,
+            *,
+            encoding: None = None,
+            **kwargs) -> StorageStreamDownloader[bytes]:
+        ...
+
     @distributed_trace_async
-    async def download_blob(self, offset=None, length=None, **kwargs):
-        # type: (Optional[int], Optional[int], Any) -> StorageStreamDownloader
+    async def download_blob(
+            self, offset: int = None,
+            length: int = None,
+            *,
+            encoding: Optional[str] = None,
+            **kwargs) -> StorageStreamDownloader:
         """Downloads a blob to the StorageStreamDownloader. The readall() method must
         be used to read all the content or readinto() must be used to download the blob into
         a stream. Using chunks() returns an async iterator which allows the user to iterate over the content in chunks.
@@ -495,6 +517,7 @@ class BlobClient(AsyncStorageAccountHostsMixin, BlobClientBase, StorageEncryptio
         options = self._download_blob_options(
             offset=offset,
             length=length,
+            encoding=encoding,
             **kwargs)
         downloader = StorageStreamDownloader(**options)
         await downloader._setup()  # pylint: disable=protected-access

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_container_client_async.py

@@ -7,7 +7,7 @@
 
 import functools
 from typing import (  # pylint: disable=unused-import
-    Any, AnyStr, AsyncIterator, Dict, List, IO, Iterable, Optional, Union,
+    Any, AnyStr, AsyncIterator, Dict, List, IO, Iterable, Optional, overload, Union,
     TYPE_CHECKING
 )
 
@@ -30,6 +30,7 @@ from .._generated.aio import AzureBlobStorage
 from .._generated.models import SignedIdentifier
 from .._container_client import ContainerClient as ContainerClientBase, _get_blob_name
 from .._deserialize import deserialize_container_properties
+from ._download_async import StorageStreamDownloader
 from .._encryption import StorageEncryptionMixin
 from .._models import ContainerProperties, BlobType, BlobProperties, FilteredBlob
 from .._serialize import get_modify_conditions, get_container_cpk_scope_info, get_api_version, get_access_conditions
@@ -40,7 +41,6 @@ from ._models import FilteredBlobPaged
 
 if TYPE_CHECKING:
     from datetime import datetime
-    from ._download_async import StorageStreamDownloader
     from .._models import ( # pylint: disable=unused-import
         AccessPolicy,
         StandardBlobTier,
@@ -928,9 +928,34 @@ class ContainerClient(AsyncStorageAccountHostsMixin, ContainerClientBase, Storag
             timeout=timeout,
             **kwargs)
 
+    @overload
+    async def download_blob(
+            self, blob: Union[str, BlobProperties],
+            offset: int = None,
+            length: int = None,
+            *,
+            encoding: str,
+            **kwargs) -> StorageStreamDownloader[str]:
+        ...
+
+    @overload
+    async def download_blob(
+            self, blob: Union[str, BlobProperties],
+            offset: int = None,
+            length: int = None,
+            *,
+            encoding: None = None,
+            **kwargs) -> StorageStreamDownloader[bytes]:
+        ...
+
     @distributed_trace_async
-    async def download_blob(self, blob, offset=None, length=None, **kwargs):
-        # type: (Union[str, BlobProperties], Optional[int], Optional[int], Any) -> StorageStreamDownloader
+    async def download_blob(
+            self, blob: Union[str, BlobProperties],
+            offset: int = None,
+            length: int = None,
+            *,
+            encoding: Optional[str] = None,
+            **kwargs) -> StorageStreamDownloader:
         """Downloads a blob to the StorageStreamDownloader. The readall() method must
         be used to read all the content or readinto() must be used to download the blob into
         a stream. Using chunks() returns an async iterator which allows the user to iterate over the content in chunks.
@@ -1014,6 +1039,7 @@ class ContainerClient(AsyncStorageAccountHostsMixin, ContainerClientBase, Storag
         return await blob_client.download_blob(
             offset=offset,
             length=length,
+            encoding=encoding,
             **kwargs)
 
     @distributed_trace_async

sdk/storage/azure-storage-blob/azure/storage/blob/aio/_download_async.py

@@ -9,7 +9,7 @@ import sys
 import warnings
 from io import BytesIO
 from itertools import islice
-from typing import AsyncIterator
+from typing import AsyncIterator, Generic, TypeVar
 
 import asyncio
 from aiohttp import ClientPayloadError
@@ -26,6 +26,8 @@ from .._encryption import (
     parse_encryption_data
 )
 
+T = TypeVar('T', bytes, str)
+
 
 async def process_content(data, start_offset, end_offset, encryption):
     if data is None:
@@ -189,7 +191,7 @@ class _AsyncChunkIterator(object):
         return chunk_data
 
 
-class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attributes
+class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-attributes
     """A streaming object to download from Azure Storage.
 
     :ivar str name:
@@ -205,18 +207,19 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
     """
 
     def __init__(
-            self,
-            clients=None,
-            config=None,
-            start_range=None,
-            end_range=None,
-            validate_content=None,
-            encryption_options=None,
-            max_concurrency=1,
-            name=None,
-            container=None,
-            encoding=None,
-            **kwargs
+        self,
+        clients=None,
+        config=None,
+        start_range=None,
+        end_range=None,
+        validate_content=None,
+        encryption_options=None,
+        max_concurrency=1,
+        name=None,
+        container=None,
+        encoding=None,
+        download_cls=None,
+        **kwargs
     ):
         self.name = name
         self.container = container
@@ -244,6 +247,10 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
         self._initial_range = None
         self._initial_offset = None
 
+        # The cls is passed in via download_cls to avoid conflicting arg name with Generic.__new__
+        # but needs to be changed to cls in the request options.
+        self._request_options['cls'] = download_cls
+
         # The service only provides transactional MD5s for chunks under 4MB.
         # If validate_content is on, get only self.MAX_CHUNK_GET_SIZE for the first
         # chunk so a transactional MD5 can be retrieved.
@@ -450,10 +457,11 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
             chunk_size=self._config.max_chunk_get_size)
 
     async def readall(self):
+        # type: () -> T
         """Download the contents of this blob.
 
         This operation is blocking until all data is downloaded.
-        :rtype: bytes or str
+        :rtype: T
         """
         stream = BytesIO()
         await self.readinto(stream)

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/_download.py

@@ -39,10 +39,11 @@ class StorageStreamDownloader(object):
         return self._downloader.chunks()
 
     def readall(self):
+        # type: () -> bytes
         """Download the contents of this file.
 
         This operation is blocking until all data is downloaded.
-        :rtype: bytes or str
+        :rtype: bytes
         """
         return self._downloader.readall()
 

sdk/storage/azure-storage-file-datalake/azure/storage/filedatalake/aio/_download_async.py

@@ -39,10 +39,11 @@ class StorageStreamDownloader(object):
         return self._downloader.chunks()
 
     async def readall(self):
+        # type: () -> bytes
         """Download the contents of this file.
 
         This operation is blocking until all data is downloaded.
-        :rtype: bytes or str
+        :rtype: bytes
         """
         return await self._downloader.readall()
 

sdk/storage/azure-storage-file-share/azure/storage/fileshare/_download.py

@@ -381,10 +381,11 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
             chunk_size=self._config.max_chunk_get_size)
 
     def readall(self):
+        # type: () -> bytes
         """Download the contents of this file.
 
         This operation is blocking until all data is downloaded.
-        :rtype: bytes or str
+        :rtype: bytes
         """
         stream = BytesIO()
         self.readinto(stream)

sdk/storage/azure-storage-file-share/azure/storage/fileshare/aio/_download_async.py

@@ -333,10 +333,11 @@ class StorageStreamDownloader(object):  # pylint: disable=too-many-instance-attr
         )
 
     async def readall(self):
+        # type: () -> bytes
         """Download the contents of this file.
 
         This operation is blocking until all data is downloaded.
-        :rtype: bytes or str
+        :rtype: bytes
         """
         stream = BytesIO()
         await self.readinto(stream)