[Storage] [Typing] [File Datalake] azure-storage-file-datalake (#39691)

Author: weirongw23-msft

sdk/storage/azure-storage-blob/azure/storage/blob/_shared/base_client.py

@@ -65,7 +65,10 @@
 
 
 class StorageAccountHostsMixin(object):
+
     _client: Any
+    _hosts: Dict[str, str]
+
     def __init__(
         self,
         parsed_url: Any,
@@ -74,7 +77,7 @@ def __init__(
         **kwargs: Any
     ) -> None:
         self._location_mode = kwargs.get("_location_mode", LocationMode.PRIMARY)
-        self._hosts = kwargs.get("_hosts")
+        self._hosts = kwargs.get("_hosts", {})
         self.scheme = parsed_url.scheme
         self._is_localhost = False
 
@@ -93,7 +96,7 @@ def __init__(
         if self.scheme.lower() != "https" and hasattr(self.credential, "get_token"):
             raise ValueError("Token credential is only supported with HTTPS.")
 
-        secondary_hostname = None
+        secondary_hostname = ""
         if hasattr(self.credential, "account_name"):
             self.account_name = self.credential.account_name
             secondary_hostname = f"{self.credential.account_name}-secondary.{service_name}.{SERVICE_HOST_BASE}"

sdk/storage/azure-storage-file-datalake/azure/__init__.py

@@ -1 +1 @@
-__path__ = __import__('pkgutil').extend_path(__path__, __name__)  # type: str
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)  # type: ignore

sdk/storage/azure-storage-file-datalake/azure/storage/__init__.py

@@ -1 +1 @@
-__path__ = __import__('pkgutil').extend_path(__path__, __name__)  # type: str
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)  # type: ignore

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

@@ -7,11 +7,10 @@
 
 import functools
 from typing import (
-    Any, Dict, Optional, Union,
+    Any, cast, Dict, Optional, Union,
     TYPE_CHECKING
 )
 from urllib.parse import quote, unquote
-
 from typing_extensions import Self
 
 from azure.core.paging import ItemPaged
@@ -23,7 +22,8 @@
 from ._list_paths_helper import PathPropertiesPaged
 from ._models import DirectoryProperties, FileProperties
 from ._path_client import PathClient
-from ._shared.base_client import TransportWrapper, parse_connection_str
+from ._path_client_helpers import _parse_rename_path
+from ._shared.base_client import parse_connection_str, TransportWrapper
 
 if TYPE_CHECKING:
     from azure.core.credentials import AzureNamedKeyCredential, AzureSasCredential, TokenCredential
@@ -37,12 +37,6 @@ class DataLakeDirectoryClient(PathClient):
     For operations relating to a specific subdirectory or file under the directory, a directory client or file client
     can be retrieved using the :func:`~get_sub_directory_client` or :func:`~get_file_client` functions.
 
-    :ivar str url:
-        The full endpoint URL to the file system, including SAS token if used.
-    :ivar str primary_endpoint:
-        The full primary endpoint URL.
-    :ivar str primary_hostname:
-        The hostname of the primary endpoint.
     :param str account_url:
         The URI to the storage account.
     :param file_system_name:
@@ -64,7 +58,7 @@ class DataLakeDirectoryClient(PathClient):
         ~azure.core.credentials.AzureNamedKeyCredential or
         ~azure.core.credentials.AzureSasCredential or
         ~azure.core.credentials.TokenCredential or
-        str or dict[str, str] or None
+        str or Dict[str, str] or None
     :keyword str api_version:
         The Storage API version to use for requests. Default value is the most recent service version that is
         compatible with the current SDK. Setting to an older version may result in reduced feature compatibility.
@@ -81,6 +75,14 @@ class DataLakeDirectoryClient(PathClient):
             :dedent: 4
             :caption: Creating the DataLakeServiceClient from connection string.
     """
+
+    url: str
+    """The full endpoint URL to the file system, including SAS token if used."""
+    primary_endpoint: str
+    """The full primary endpoint URL."""
+    primary_hostname: str
+    """The hostname of the primary endpoint."""
+
     def __init__(
         self, account_url: str,
         file_system_name: str,
@@ -93,12 +95,12 @@ def __init__(
 
     @classmethod
     def from_connection_string(
-            cls, conn_str: str,
-            file_system_name: str,
-            directory_name: str,
-            credential: Optional[Union[str, Dict[str, str], "AzureNamedKeyCredential", "AzureSasCredential", "TokenCredential"]] = None,  # pylint: disable=line-too-long
-            **kwargs: Any
-        ) -> Self:
+        cls, conn_str: str,
+        file_system_name: str,
+        directory_name: str,
+        credential: Optional[Union[str, Dict[str, str], "AzureNamedKeyCredential", "AzureSasCredential", "TokenCredential"]] = None,  # pylint: disable=line-too-long
+        **kwargs: Any
+    ) -> Self:
         """
         Create DataLakeDirectoryClient from a Connection String.
 
@@ -120,7 +122,7 @@ def from_connection_string(
             ~azure.core.credentials.AzureNamedKeyCredential or
             ~azure.core.credentials.AzureSasCredential or
             ~azure.core.credentials.TokenCredential or
-            str or dict[str, str] or None
+            str or Dict[str, str] or None
         :param directory_name:
             The name of directory to interact with. The directory is under file system.
         :type directory_name: str
@@ -136,15 +138,16 @@ def from_connection_string(
             credential=credential, **kwargs)
 
     @distributed_trace
-    def create_directory(self, metadata=None,  # type: Optional[Dict[str, str]]
-                         **kwargs):
-        # type: (...) -> Dict[str, Union[str, datetime]]
+    def create_directory(
+        self, metadata: Optional[Dict[str, str]] = None,
+        **kwargs: Any
+    ) -> Dict[str, Union[str, "datetime"]]:
         """
         Create a new directory.
 
         :param metadata:
             Name-value pairs associated with the file as metadata.
-        :type metadata: dict(str, str)
+        :type metadata: Dict[str, str]
         :keyword ~azure.storage.filedatalake.ContentSettings content_settings:
             ContentSettings object used to set path properties.
         :keyword lease:
@@ -210,7 +213,7 @@ def create_directory(self, metadata=None,  # type: Optional[Dict[str, str]]
             see `here <https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/storage/azure-storage-file-datalake
             #other-client--per-operation-configuration>`_.
         :return: A dictionary of response headers.
-        :rtype: dict[str, Union[str, datetime]]
+        :rtype: Dict[str, Union[str, ~datetime.datetime]]
 
         .. admonition:: Example:
 
@@ -224,8 +227,7 @@ def create_directory(self, metadata=None,  # type: Optional[Dict[str, str]]
         return self._create('directory', metadata=metadata, **kwargs)
 
     @distributed_trace
-    def delete_directory(self, **kwargs):
-        # type: (...) -> None
+    def delete_directory(self, **kwargs: Any) -> Dict[str, Any]:  # pylint: disable=delete-operation-wrong-return-type
         """
         Marks the specified directory for deletion.
 
@@ -256,8 +258,8 @@ def delete_directory(self, **kwargs):
             This value is not tracked or validated on the client. To configure client-side network timesouts
             see `here <https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/storage/azure-storage-file-datalake
             #other-client--per-operation-configuration>`_.
-        :returns: None.
-        :rtype: None
+        :returns: A dictionary of response headers.
+        :rtype: Dict[str, Any]
 
         .. admonition:: Example:
 
@@ -271,8 +273,7 @@ def delete_directory(self, **kwargs):
         return self._delete(recursive=True, **kwargs)
 
     @distributed_trace
-    def get_directory_properties(self, **kwargs):
-        # type: (**Any) -> DirectoryProperties
+    def get_directory_properties(self, **kwargs: Any) -> DirectoryProperties:
         """Returns all user-defined metadata, standard HTTP properties, and
         system properties for the directory. It does not return the content of the directory.
 
@@ -333,11 +334,10 @@ def get_directory_properties(self, **kwargs):
             headers = kwargs.pop('headers', {})
             headers['x-ms-upn'] = str(upn)
             kwargs['headers'] = headers
-        return self._get_path_properties(cls=deserialize_dir_properties, **kwargs)
+        return cast(DirectoryProperties, self._get_path_properties(cls=deserialize_dir_properties, **kwargs))
 
     @distributed_trace
-    def exists(self, **kwargs):
-        # type: (**Any) -> bool
+    def exists(self, **kwargs: Any) -> bool:
         """
         Returns True if a directory exists and returns False otherwise.
 
@@ -353,8 +353,7 @@ def exists(self, **kwargs):
         return self._exists(**kwargs)
 
     @distributed_trace
-    def rename_directory(self, new_name, **kwargs):
-        # type: (str, **Any) -> DataLakeDirectoryClient
+    def rename_directory(self, new_name: str, **kwargs: Any) -> "DataLakeDirectoryClient":
         """
         Rename the source directory.
 
@@ -422,7 +421,8 @@ def rename_directory(self, new_name, **kwargs):
                 :dedent: 4
                 :caption: Rename the source directory.
         """
-        new_file_system, new_path, new_dir_sas = self._parse_rename_path(new_name)
+        new_file_system, new_path, new_dir_sas = _parse_rename_path(
+            new_name, self.file_system_name, self._query_str, self._raw_credential)
 
         new_directory_client = DataLakeDirectoryClient(
             f"{self.scheme}://{self.primary_hostname}", new_file_system, directory_name=new_path,
@@ -433,10 +433,11 @@ def rename_directory(self, new_name, **kwargs):
         return new_directory_client
 
     @distributed_trace
-    def create_sub_directory(self, sub_directory,  # type: Union[DirectoryProperties, str]
-                             metadata=None,  # type: Optional[Dict[str, str]]
-                             **kwargs):
-        # type: (...) -> DataLakeDirectoryClient
+    def create_sub_directory(
+        self, sub_directory: Union[DirectoryProperties, str],
+        metadata: Optional[Dict[str, str]] = None,
+        **kwargs: Any
+    ) -> "DataLakeDirectoryClient":
         """
         Create a subdirectory and return the subdirectory client to be interacted with.
 
@@ -446,7 +447,7 @@ def create_sub_directory(self, sub_directory,  # type: Union[DirectoryProperties
         :type sub_directory: str or ~azure.storage.filedatalake.DirectoryProperties
         :param metadata:
             Name-value pairs associated with the file as metadata.
-        :type metadata: dict(str, str)
+        :type metadata: Dict[str, str]
         :keyword ~azure.storage.filedatalake.ContentSettings content_settings:
             ContentSettings object used to set path properties.
         :keyword lease:
@@ -519,9 +520,10 @@ def create_sub_directory(self, sub_directory,  # type: Union[DirectoryProperties
         return subdir
 
     @distributed_trace
-    def delete_sub_directory(self, sub_directory,  # type: Union[DirectoryProperties, str]
-                             **kwargs):
-        # type: (...) -> DataLakeDirectoryClient
+    def delete_sub_directory(  # pylint: disable=delete-operation-wrong-return-type
+        self, sub_directory: Union[DirectoryProperties, str],
+        **kwargs: Any
+    ) -> "DataLakeDirectoryClient":
         """
         Marks the specified subdirectory for deletion.
 
@@ -564,9 +566,7 @@ def delete_sub_directory(self, sub_directory,  # type: Union[DirectoryProperties
         return subdir
 
     @distributed_trace
-    def create_file(self, file,  # type: Union[FileProperties, str]
-                    **kwargs):
-        # type: (...) -> DataLakeFileClient
+    def create_file(self, file: Union[FileProperties, str], **kwargs: Any) -> DataLakeFileClient:
         """
         Create a new file and return the file client to be interacted with.
 
@@ -578,7 +578,7 @@ def create_file(self, file,  # type: Union[FileProperties, str]
             ContentSettings object used to set path properties.
         :keyword metadata:
             Name-value pairs associated with the file as metadata.
-        :type metadata: dict(str, str)
+        :type metadata: Dict[str, str]
         :keyword lease:
             Required if the file has an active lease. Value can be a DataLakeLeaseClient object
             or the lease ID as a string.
@@ -688,7 +688,6 @@ def get_paths(
         :returns: An iterable (auto-paging) response of PathProperties.
         :rtype: ~azure.core.paging.ItemPaged[~azure.storage.filedatalake.PathProperties]
         """
-        timeout = kwargs.pop('timeout', None)
         hostname = self._hosts[self._location_mode]
         url = f"{self.scheme}://{hostname}/{quote(self.file_system_name)}"
         client = self._build_generated_client(url)
@@ -702,9 +701,7 @@ def get_paths(
             command, recursive, path=self.path_name, max_results=max_results,
             upn=upn, page_iterator_class=PathPropertiesPaged, **kwargs)
 
-    def get_file_client(self, file  # type: Union[FileProperties, str]
-                        ):
-        # type: (...) -> DataLakeFileClient
+    def get_file_client(self, file: Union[FileProperties, str]) -> DataLakeFileClient:
         """Get a client to interact with the specified file.
 
         The file need not already exist.
@@ -716,9 +713,9 @@ def get_file_client(self, file  # type: Union[FileProperties, str]
         :returns: A DataLakeFileClient.
         :rtype: ~azure.storage.filedatalake.DataLakeFileClient
         """
-        try:
+        if isinstance(file, FileProperties):
             file_path = file.get('name')
-        except AttributeError:
+        else:
             file_path = self.path_name + '/' + str(file)
 
         _pipeline = Pipeline(
@@ -730,9 +727,7 @@ def get_file_client(self, file  # type: Union[FileProperties, str]
             api_version=self.api_version,
             _hosts=self._hosts, _configuration=self._config, _pipeline=_pipeline)
 
-    def get_sub_directory_client(self, sub_directory  # type: Union[DirectoryProperties, str]
-                                 ):
-        # type: (...) -> DataLakeDirectoryClient
+    def get_sub_directory_client(self, sub_directory: Union[DirectoryProperties, str]) -> "DataLakeDirectoryClient":
         """Get a client to interact with the specified subdirectory of the current directory.
 
         The sub subdirectory need not already exist.
@@ -744,14 +739,14 @@ def get_sub_directory_client(self, sub_directory  # type: Union[DirectoryPropert
         :returns: A DataLakeDirectoryClient.
         :rtype: ~azure.storage.filedatalake.DataLakeDirectoryClient
         """
-        try:
+        if isinstance(sub_directory, DirectoryProperties):
             subdir_path = sub_directory.get('name')
-        except AttributeError:
+        else:
             subdir_path = self.path_name + '/' + str(sub_directory)
 
         _pipeline = Pipeline(
-            transport=TransportWrapper(self._pipeline._transport), # pylint: disable = protected-access
-            policies=self._pipeline._impl_policies # pylint: disable = protected-access
+            transport=TransportWrapper(self._pipeline._transport), # pylint: disable=protected-access
+            policies=self._pipeline._impl_policies # pylint: disable=protected-access
         )
         return DataLakeDirectoryClient(
             self.url, self.file_system_name, directory_name=subdir_path, credential=self._raw_credential,