azure-storage-blob 12.21.0__py3-none-any.whl → 12.22.0__py3-none-any.whl

azure/storage/blob/_download.py

@@ -9,13 +9,17 @@ import threading
 import time
 import warnings
 from io import BytesIO, StringIO
-from typing import Generic, IO, Iterator, Optional, overload, TypeVar, Union
+from typing import (
+    Any, Callable, cast, Dict, Generator,
+    Generic, IO, Iterator, List, Optional,
+    overload, Tuple, TypeVar, Union, TYPE_CHECKING
+)
 
 from azure.core.exceptions import DecodeError, HttpResponseError, IncompleteReadError
 from azure.core.tracing.common import with_current_context
 
 from ._shared.request_handlers import validate_and_format_range_headers
-from ._shared.response_handlers import process_storage_error, parse_length_from_content_range
+from ._shared.response_handlers import parse_length_from_content_range, process_storage_error
 from ._deserialize import deserialize_blob_properties, get_page_ranges_result
 from ._encryption import (
     adjust_blob_size_for_encryption,
@@ -25,10 +29,25 @@ from ._encryption import (
     parse_encryption_data
 )
 
+if TYPE_CHECKING:
+    from codecs import IncrementalDecoder
+    from ._encryption import _EncryptionData
+    from ._generated import AzureBlobStorage
+    from ._generated.operations import BlobOperations
+    from ._models import BlobProperties
+    from ._shared.models import StorageConfiguration
+
+
 T = TypeVar('T', bytes, str)
 
 
-def process_range_and_offset(start_range, end_range, length, encryption_options, encryption_data):
+def process_range_and_offset(
+    start_range: int,
+    end_range: int,
+    length: Optional[int],
+    encryption_options: Dict[str, Any],
+    encryption_data: Optional["_EncryptionData"]
+) -> Tuple[Tuple[int, int], Tuple[int, int]]:
     start_offset, end_offset = 0, 0
     if encryption_options.get("key") is not None or encryption_options.get("resolver") is not None:
         return get_adjusted_download_range_and_offset(
@@ -40,7 +59,7 @@ def process_range_and_offset(start_range, end_range, length, encryption_options,
     return (start_range, end_range), (start_offset, end_offset)
 
 
-def process_content(data, start_offset, end_offset, encryption):
+def process_content(data: Any, start_offset: int, end_offset: int, encryption: Dict[str, Any]) -> bytes:
     if data is None:
         raise ValueError("Response cannot be None.")
 
@@ -49,7 +68,7 @@ def process_content(data, start_offset, end_offset, encryption):
     if content and encryption.get("key") is not None or encryption.get("resolver") is not None:
         try:
             return decrypt_blob(
-                encryption.get("required"),
+                encryption.get("required") or False,
                 encryption.get("key"),
                 encryption.get("resolver"),
                 content,
@@ -65,21 +84,21 @@ def process_content(data, start_offset, end_offset, encryption):
 class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
-        client=None,
-        non_empty_ranges=None,
-        total_size=None,
-        chunk_size=None,
-        current_progress=None,
-        start_range=None,
-        end_range=None,
-        stream=None,
-        parallel=None,
-        validate_content=None,
-        encryption_options=None,
-        encryption_data=None,
-        progress_hook=None,
-        **kwargs
-    ):
+        client: "BlobOperations",
+        total_size: int,
+        chunk_size: int,
+        current_progress: int,
+        start_range: int,
+        end_range: int,
+        validate_content: bool,
+        encryption_options: Dict[str, Any],
+        encryption_data: Optional["_EncryptionData"] = None,
+        stream: Any = None,
+        parallel: Optional[int] = None,
+        non_empty_ranges: Optional[List[Dict[str, Any]]] = None,
+        progress_hook: Optional[Callable[[int, Optional[int]], None]] = None,
+        **kwargs: Any
+    ) -> None:
         self.client = client
         self.non_empty_ranges = non_empty_ranges
 
@@ -110,20 +129,20 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
         self.validate_content = validate_content
         self.request_options = kwargs
 
-    def _calculate_range(self, chunk_start):
+    def _calculate_range(self, chunk_start: int) -> Tuple[int, int]:
         if chunk_start + self.chunk_size > self.end_index:
             chunk_end = self.end_index
         else:
             chunk_end = chunk_start + self.chunk_size
         return chunk_start, chunk_end
 
-    def get_chunk_offsets(self):
+    def get_chunk_offsets(self) -> Generator[int, None, None]:
         index = self.start_index
         while index < self.end_index:
             yield index
             index += self.chunk_size
 
-    def process_chunk(self, chunk_start):
+    def process_chunk(self, chunk_start: int) -> None:
         chunk_start, chunk_end = self._calculate_range(chunk_start)
         chunk_data, _ = self._download_chunk(chunk_start, chunk_end - 1)
         length = chunk_end - chunk_start
@@ -131,11 +150,11 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
             self._write_to_stream(chunk_data, chunk_start)
             self._update_progress(length)
 
-    def yield_chunk(self, chunk_start):
+    def yield_chunk(self, chunk_start: int) -> Tuple[bytes, int]:
         chunk_start, chunk_end = self._calculate_range(chunk_start)
         return self._download_chunk(chunk_start, chunk_end - 1)
 
-    def _update_progress(self, length):
+    def _update_progress(self, length: int) -> None:
         if self.progress_lock:
             with self.progress_lock:  # pylint: disable=not-context-manager
                 self.progress_total += length
@@ -145,7 +164,7 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
         if self.progress_hook:
             self.progress_hook(self.progress_total, self.total_size)
 
-    def _write_to_stream(self, chunk_data, chunk_start):
+    def _write_to_stream(self, chunk_data: bytes, chunk_start: int) -> None:
         if self.stream_lock:
             with self.stream_lock:  # pylint: disable=not-context-manager
                 self.stream.seek(self.stream_start + (chunk_start - self.start_index))
@@ -153,7 +172,7 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
         else:
             self.stream.write(chunk_data)
 
-    def _do_optimize(self, given_range_start, given_range_end):
+    def _do_optimize(self, given_range_start: int, given_range_end: int) -> bool:
         # If we have no page range list stored, then assume there's data everywhere for that page blob
         # or it's a block blob or append blob
         if self.non_empty_ranges is None:
@@ -178,7 +197,9 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
         # Went through all src_ranges, but nothing overlapped. Optimization will be applied.
         return True
 
-    def _download_chunk(self, chunk_start, chunk_end):
+    def _download_chunk(self, chunk_start: int, chunk_end: int) -> Tuple[bytes, int]:
+        if self.encryption_options is None:
+            raise ValueError("Required argument is missing: encryption_options")
         download_range, offset = process_range_and_offset(
             chunk_start, chunk_end, chunk_end, self.encryption_options, self.encryption_data
         )
@@ -198,6 +219,7 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
             retry_active = True
             retry_total = 3
             while retry_active:
+                response: Any = None
                 try:
                     _, response = self.client.download(
                         range=range_header,
@@ -229,24 +251,24 @@ class _ChunkDownloader(object):  # pylint: disable=too-many-instance-attributes
 
 
 class _ChunkIterator(object):
-    """Async iterator for chunks in blob download stream."""
+    """Iterator for chunks in blob download stream."""
 
-    def __init__(self, size, content, downloader, chunk_size):
+    def __init__(self, size: int, content: bytes, downloader: Optional[_ChunkDownloader], chunk_size: int) -> None:
         self.size = size
         self._chunk_size = chunk_size
         self._current_content = content
         self._iter_downloader = downloader
-        self._iter_chunks = None
+        self._iter_chunks: Optional[Generator[int, None, None]] = None
         self._complete = size == 0
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.size
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[bytes]:
         return self
 
     # Iterate through responses.
-    def __next__(self):
+    def __next__(self) -> bytes:
         if self._complete:
             raise StopIteration("Download complete")
         if not self._iter_downloader:
@@ -264,8 +286,8 @@ class _ChunkIterator(object):
             return self._get_chunk_data()
 
         try:
-            chunk = next(self._iter_chunks)
-            self._current_content += self._iter_downloader.yield_chunk(chunk)[0]
+            next_chunk = next(self._iter_chunks)
+            self._current_content += self._iter_downloader.yield_chunk(next_chunk)[0]
         except StopIteration as e:
             self._complete = True
             if self._current_content:
@@ -278,46 +300,46 @@ class _ChunkIterator(object):
 
     next = __next__  # Python 2 compatibility.
 
-    def _get_chunk_data(self):
+    def _get_chunk_data(self) -> bytes:
         chunk_data = self._current_content[: self._chunk_size]
         self._current_content = self._current_content[self._chunk_size:]
         return chunk_data
 
 
 class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-attributes
-    """A streaming object to download from Azure Storage.
-
-    :ivar str name:
-        The name of the blob being downloaded.
-    :ivar str container:
-        The name of the container where the blob is.
-    :ivar ~azure.storage.blob.BlobProperties properties:
-        The properties of the blob being downloaded. If only a range of the data is being
-        downloaded, this will be reflected in the properties.
-    :ivar int size:
-        The size of the total data in the stream. This will be the byte range if specified,
-        otherwise the total size of the blob.
     """
+    A streaming object to download from Azure Storage.
+    """
+
+    name: str
+    """The name of the blob being downloaded."""
+    container: str
+    """The name of the container where the blob is."""
+    properties: "BlobProperties"
+    """The properties of the blob being downloaded. If only a range of the data is being
+    downloaded, this will be reflected in the properties."""
+    size: int
+    """The size of the total data in the stream. This will be the byte range if specified,
+    otherwise the total size of the blob."""
 
     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,
-        download_cls=None,
-        **kwargs
-    ):
+        clients: "AzureBlobStorage" = None,  # type: ignore [assignment]
+        config: "StorageConfiguration" = None,  # type: ignore [assignment]
+        start_range: Optional[int] = None,
+        end_range: Optional[int] = None,
+        validate_content: bool = None,  # type: ignore [assignment]
+        encryption_options: Dict[str, Any] = None,  # type: ignore [assignment]
+        max_concurrency: int = 1,
+        name: str = None,  # type: ignore [assignment]
+        container: str = None,  # type: ignore [assignment]
+        encoding: Optional[str] = None,
+        download_cls: Optional[Callable] = None,
+        **kwargs: Any
+    ) -> None:
         self.name = name
         self.container = container
-        self.properties = None
-        self.size = None
+        self.size = 0
 
         self._clients = clients
         self._config = config
@@ -331,10 +353,10 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         self._request_options = kwargs
         self._response = None
         self._location_mode = None
-        self._current_content = b''
+        self._current_content: Union[str, bytes] = b''
         self._file_size = 0
         self._non_empty_ranges = None
-        self._encryption_data = None
+        self._encryption_data: Optional["_EncryptionData"] = None
 
         # The content download offset, after any processing (decryption), in bytes
         self._download_offset = 0
@@ -346,7 +368,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         self._current_content_offset = 0
 
         self._text_mode: Optional[bool] = None
-        self._decoder = None
+        self._decoder: Optional["IncrementalDecoder"] = None
         # Whether the current content is the first chunk of download content or not
         self._first_chunk = True
         self._download_start = self._start_range or 0
@@ -379,7 +401,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         )
 
         self._response = self._initial_request()
-        self.properties = self._response.properties
+        self.properties = cast("BlobProperties", self._response.properties)
         self.properties.name = self.name
         self.properties.container = self.container
 
@@ -392,18 +414,18 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         # Overwrite the content MD5 as it is the MD5 for the last range instead
         # of the stored MD5
         # TODO: Set to the stored MD5 when the service returns this
-        self.properties.content_md5 = None
+        self.properties.content_md5 = None  # type: ignore [attr-defined]
 
     def __len__(self):
         return self.size
 
-    def _get_encryption_data_request(self):
+    def _get_encryption_data_request(self) -> None:
         # Save current request cls
         download_cls = self._request_options.pop('cls', None)
         # Adjust cls for get_properties
         self._request_options['cls'] = deserialize_blob_properties
 
-        properties = self._clients.blob.get_properties(**self._request_options)
+        properties = cast("BlobProperties", self._clients.blob.get_properties(**self._request_options))
         # This will return None if there is no encryption metadata or there are parsing errors.
         # That is acceptable here, the proper error will be caught and surfaced when attempting
         # to decrypt the blob.
@@ -431,14 +453,14 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         retry_total = 3
         while retry_active:
             try:
-                location_mode, response = self._clients.blob.download(
+                location_mode, response = cast(Tuple[Optional[str], Any], self._clients.blob.download(
                     range=range_header,
                     range_get_content_md5=range_validation,
                     validate_content=self._validate_content,
                     data_stream_total=None,
                     download_stream_current=0,
                     **self._request_options
-                )
+                ))
 
                 # Check the location we read from to ensure we use the same one
                 # for subsequent requests.
@@ -452,7 +474,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
                 # Remove any extra encryption data size from blob size
                 self._file_size = adjust_blob_size_for_encryption(self._file_size, self._encryption_data)
 
-                if self._end_range is not None:
+                if self._end_range is not None and self._start_range is not None:
                     # Use the end range index unless it is over the end of the file
                     self.size = min(self._file_size - self._start_range, self._end_range - self._start_range + 1)
                 elif self._start_range is not None:
@@ -517,8 +539,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
 
         return response
 
-    def chunks(self):
-        # type: () -> Iterator[bytes]
+    def chunks(self) -> Iterator[bytes]:
         """
         Iterate over chunks in the download stream. Note, the iterator returned will
         iterate over the entire download content, regardless of any data that was
@@ -573,7 +594,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         initial_content = self._current_content if self._first_chunk else b''
         return _ChunkIterator(
             size=self.size,
-            content=initial_content,
+            content=cast(bytes, initial_content),
             downloader=iter_downloader,
             chunk_size=self._config.max_chunk_get_size)
 
@@ -622,12 +643,13 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         # Empty blob or already read to the end
         if (size == 0 or chars == 0 or
                 (self._download_complete and self._current_content_offset >= len(self._current_content))):
-            return b'' if not self._encoding else ''
+            return b'' if not self._encoding else ''  # type: ignore [return-value]
 
-        if not self._text_mode and chars is not None:
+        if not self._text_mode and chars is not None and self._encoding is not None:
             self._text_mode = True
             self._decoder = codecs.getincrementaldecoder(self._encoding)('strict')
-            self._current_content = self._decoder.decode(self._current_content, final=self._download_complete)
+            self._current_content = self._decoder.decode(
+                cast(bytes, self._current_content), final=self._download_complete)
         elif self._text_mode is None:
             self._text_mode = False
 
@@ -644,7 +666,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         # Start by reading from current_content
         start = self._current_content_offset
         length = min(len(self._current_content) - self._current_content_offset, size - count)
-        read = output_stream.write(self._current_content[start:start + length])
+        read = output_stream.write(self._current_content[start:start + length])  # type: ignore [arg-type]
 
         count += read
         self._current_content_offset += read
@@ -691,8 +713,8 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
                             downloader.get_chunk_offsets()
                         ))
                 else:
-                    for chunk in chunks_iter:
-                        downloader.process_chunk(chunk)
+                    for next_chunk in chunks_iter:
+                        downloader.process_chunk(next_chunk)
 
                 self._complete_read()
 
@@ -701,13 +723,15 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
                     chunk_data, content_length = downloader.yield_chunk(chunk)
                     self._download_offset += len(chunk_data)
                     self._raw_download_offset += content_length
-                    self._current_content = self._decoder.decode(
-                        chunk_data, final=self._download_complete) if self._text_mode else chunk_data
+                    if self._text_mode and self._decoder is not None:
+                        self._current_content = self._decoder.decode(chunk_data, final=self._download_complete)
+                    else:
+                        self._current_content = chunk_data
 
                     if remaining < len(self._current_content):
-                        read = output_stream.write(self._current_content[:remaining])
+                        read = output_stream.write(self._current_content[:remaining])  # type: ignore [arg-type]
                     else:
-                        read = output_stream.write(self._current_content)
+                        read = output_stream.write(self._current_content)  # type: ignore [arg-type]
 
                     self._current_content_offset = read
                     self._read_offset += read
@@ -718,7 +742,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         if not self._text_mode and self._encoding:
             try:
                 # This is technically incorrect to do, but we have it for backwards compatibility.
-                data = data.decode(self._encoding)
+                data = cast(bytes, data).decode(self._encoding)
             except UnicodeDecodeError:
                 warnings.warn(
                     "Encountered a decoding error while decoding blob data from a partial read. "
@@ -726,7 +750,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
                 )
                 raise
 
-        return data
+        return data  # type: ignore [return-value]
 
     def readall(self) -> T:
         """
@@ -774,7 +798,7 @@ class StorageStreamDownloader(Generic[T]):  # pylint: disable=too-many-instance-
         # Write the current content to the user stream
         current_remaining = len(self._current_content) - self._current_content_offset
         start = self._current_content_offset
-        count = stream.write(self._current_content[start:start + current_remaining])
+        count = stream.write(cast(bytes, self._current_content[start:start + current_remaining]))
 
         self._current_content_offset += count
         self._read_offset += count

azure/storage/blob/_encryption.py

@@ -47,6 +47,10 @@ _GCM_TAG_LENGTH = 16
 _ERROR_OBJECT_INVALID = \
     '{0} does not define a complete interface. Value of {1} is either missing or invalid.'
 
+_ERROR_UNSUPPORTED_METHOD_FOR_ENCRYPTION = (
+    'The require_encryption flag is set, but encryption is not supported'
+    ' for this method.')
+
 
 class KeyEncryptionKey(Protocol):
 
@@ -357,7 +361,7 @@ def get_adjusted_upload_size(length: int, encryption_version: str) -> int:
 def get_adjusted_download_range_and_offset(
         start: int,
         end: int,
-        length: int,
+        length: Optional[int],
         encryption_data: Optional[_EncryptionData]) -> Tuple[Tuple[int, int], Tuple[int, int]]:
     """
     Gets the new download range and offsets into the decrypted data for
@@ -374,7 +378,7 @@ def get_adjusted_download_range_and_offset(
 
     :param int start: The user-requested start index.
     :param int end: The user-requested end index.
-    :param int length: The user-requested length. Only used for V1.
+    :param Optional[int] length: The user-requested length. Only used for V1.
     :param Optional[_EncryptionData] encryption_data: The encryption data to determine version and sizes.
     :return: (new start, new end), (start offset, end offset)
     :rtype: Tuple[Tuple[int, int], Tuple[int, int]]
@@ -451,17 +455,20 @@ def parse_encryption_data(metadata: Dict[str, Any]) -> Optional[_EncryptionData]
         return None
 
 
-def adjust_blob_size_for_encryption(size: int, encryption_data: _EncryptionData) -> int:
+def adjust_blob_size_for_encryption(size: int, encryption_data: Optional[_EncryptionData]) -> int:
     """
     Adjusts the given blob size for encryption by subtracting the size of
     the encryption data (nonce + tag). This only has an affect for encryption V2.
 
     :param int size: The original blob size.
-    :param _EncryptionData encryption_data: The encryption data to determine version and sizes.
+    :param Optional[_EncryptionData] encryption_data: The encryption data to determine version and sizes.
     :return: The new blob size.
     :rtype: int
     """
-    if is_encryption_v2(encryption_data) and encryption_data.encrypted_region_info is not None:
+    if (encryption_data is not None and
+        encryption_data.encrypted_region_info is not None and
+        is_encryption_v2(encryption_data)):
+
         nonce_length = encryption_data.encrypted_region_info.nonce_length
         data_length = encryption_data.encrypted_region_info.data_length
         tag_length = encryption_data.encrypted_region_info.tag_length
@@ -836,7 +843,7 @@ def generate_blob_encryption_data(
 
 def decrypt_blob(  # pylint: disable=too-many-locals,too-many-statements
         require_encryption: bool,
-        key_encryption_key: KeyEncryptionKey,
+        key_encryption_key: Optional[KeyEncryptionKey],
         key_resolver: Optional[Callable[[str], KeyEncryptionKey]],
         content: bytes,
         start_offset: int,
@@ -848,7 +855,7 @@ def decrypt_blob(  # pylint: disable=too-many-locals,too-many-statements
 
     :param bool require_encryption:
         Whether the calling blob service requires objects to be decrypted.
-    :param KeyEncryptionKey key_encryption_key:
+    :param Optional[KeyEncryptionKey] key_encryption_key:
         The user-provided key-encryption-key. Must implement the following methods:
         wrap_key(key)
             - Wraps the specified key using an algorithm of the user's choice.