move filename determination to models.Response and file and progress bar management to clients

Author: jreiberkyle

planet/clients/data.py

@@ -20,6 +20,8 @@
 from typing import Any, AsyncIterator, Callable, Dict, List, Optional
 import uuid
 
+from tqdm.asyncio import tqdm
+
 from ..data_filter import empty_filter
 from .. import exceptions
 from ..constants import PLANET_BASE_URL
@@ -586,20 +588,40 @@ async def download_asset(self,
 
         Raises:
             planet.exceptions.APIError: On API error.
-            planet.exceptions.ClientError: If asset is not active, asset
-            description is not valid, or retry limit is exceeded.
+            planet.exceptions.ClientError: If asset is not active or asset
+            description is not valid.
         """
         try:
             location = asset['location']
         except KeyError:
             raise exceptions.ClientError(
                 'asset missing ["location"] entry. Is asset active?')
 
-        return await self._session.write(location,
-                                         filename=filename,
-                                         directory=directory,
-                                         overwrite=overwrite,
-                                         progress_bar=progress_bar)
+        response = await self._session.request(method='GET', url=location)
+        filename = filename or response.filename
+        if not filename:
+            raise exceptions.ClientError(
+                f'Could not determine filename at {location}')
+
+        dl_path = Path(directory, filename)
+        dl_path.parent.mkdir(exist_ok=True, parents=True)
+        LOGGER.info(f'Downloading {dl_path}')
+
+        try:
+            mode = 'wb' if overwrite else 'xb'
+            with open(dl_path, mode) as fp:
+                with tqdm(total=response.length,
+                          unit_scale=True,
+                          unit_divisor=1024 * 1024,
+                          unit='B',
+                          desc=str(filename),
+                          disable=not progress_bar) as progress:
+                    update = progress.update if progress_bar else LOGGER.debug
+                    await self._session.write(location, fp, update)
+        except FileExistsError:
+            LOGGER.info(f'File {dl_path} exists, not overwriting')
+
+        return dl_path
 
     @staticmethod
     def validate_checksum(asset: dict, filename: Path):

planet/clients/orders.py

@@ -15,13 +15,15 @@
 """Functionality for interacting with the orders api"""
 import asyncio
 import logging
+from pathlib import Path
 import time
 from typing import AsyncIterator, Callable, List, Optional
 import uuid
 import json
 import hashlib
 
-from pathlib import Path
+from tqdm.asyncio import tqdm
+
 from .. import exceptions
 from ..constants import PLANET_BASE_URL
 from ..http import Session
@@ -255,11 +257,31 @@ async def download_asset(self,
                 limit is exceeded.
 
         """
-        return await self._session.write(location,
-                                         filename=filename,
-                                         directory=directory,
-                                         overwrite=overwrite,
-                                         progress_bar=progress_bar)
+        response = await self._session.request(method='GET', url=location)
+        filename = filename or response.filename
+        length = response.length
+        if not filename:
+            raise exceptions.ClientError(
+                f'Could not determine filename at {location}')
+
+        dl_path = Path(directory, filename)
+        dl_path.parent.mkdir(exist_ok=True, parents=True)
+        LOGGER.info(f'Downloading {dl_path}')
+
+        try:
+            mode = 'wb' if overwrite else 'xb'
+            with open(dl_path, mode) as fp:
+                with tqdm(total=length,
+                          unit_scale=True,
+                          unit_divisor=1024 * 1024,
+                          unit='B',
+                          desc=str(filename),
+                          disable=not progress_bar) as progress:
+                    await self._session.write(location, fp, progress.update)
+        except FileExistsError:
+            LOGGER.info(f'File {dl_path} exists, not overwriting')
+
+        return dl_path
 
     async def download_order(self,
                              order_id: str,

planet/http.py

@@ -18,17 +18,11 @@
 from collections import Counter
 from http import HTTPStatus
 import logging
-import mimetypes
-from pathlib import Path
 import random
-import re
-import string
 import time
-from typing import Optional
-from urllib.parse import urlparse
+from typing import Callable, Optional
 
 import httpx
-from tqdm.asyncio import tqdm
 from typing_extensions import Literal
 
 from .auth import Auth, AuthType
@@ -332,6 +326,7 @@ async def _retry(self, func, *a, **kw):
                         LOGGER.info(f'Retrying: sleeping {wait_time}s')
                         await asyncio.sleep(wait_time)
                 else:
+                    LOGGER.info('Retrying: failed')
                     raise e
 
         self.outcomes.update(['Successful'])
@@ -399,80 +394,32 @@ async def _send(self, request, stream=False) -> httpx.Response:
 
         return http_resp
 
-    async def write(self,
-                    url: str,
-                    filename: Optional[str] = None,
-                    directory: Path = Path('.'),
-                    overwrite: bool = False,
-                    progress_bar: bool = False) -> Path:
+    async def write(self, url: str, fp, callback: Optional[Callable] = None):
         """Write data to local file with limiting and retries.
 
         Parameters:
-            url: Remote location url
-            filename: Custom name to assign to downloaded file.
-            directory: Base directory for file download. This directory will be
-                created if it does not already exist.
-            overwrite: Overwrite any existing files.
-            progress_bar: Show progress bar during download.
-
-        Returns:
-            Path to downloaded file.
+            url: Remote location url.
+            fp: Open write file pointer.
+            callback: Function that handles write progress updates.
 
         Raises:
             planet.exceptions.APIException: On API error.
-            planet.exceptions.ClientError: When retry limit is exceeded.
 
         """
 
-        async def _write():
-            async with self._client.stream('GET', url) as response:
-
-                dl_path = Path(
-                    directory,
-                    filename or _get_filename_from_response(response))
-                dl_path.parent.mkdir(exist_ok=True, parents=True)
-
-                await self._write_response(response,
-                                           dl_path,
-                                           overwrite=overwrite,
-                                           progress_bar=progress_bar)
-
-                return dl_path
-
         async def _limited_write():
             async with self._limiter:
-                dl_path = await _write()
-            return dl_path
-
-        return await self._retry(_limited_write)
-
-    async def _write_response(self,
-                              response,
-                              filename,
-                              overwrite,
-                              progress_bar):
-        total = int(response.headers["Content-Length"])
-
-        try:
-            mode = 'wb' if overwrite else 'xb'
-            with open(filename, mode) as fp:
-
-                with tqdm(total=total,
-                          unit_scale=True,
-                          unit_divisor=1024 * 1024,
-                          unit='B',
-                          desc=str(filename),
-                          disable=not progress_bar) as progress:
+                async with self._client.stream('GET', url) as response:
                     previous = response.num_bytes_downloaded
 
                     async for chunk in response.aiter_bytes():
                         fp.write(chunk)
-                        new = response.num_bytes_downloaded - previous
-                        progress.update(new - previous)
-                        previous = new
-                    progress.update()
-        except FileExistsError:
-            LOGGER.info(f'File {filename} exists, not overwriting')
+                        current = response.num_bytes_downloaded
+                        if callback is not None:
+                            callback(current - previous)
+                        previous = current
+
+        await self._retry(_limited_write)
 
     def client(self,
                name: Literal['data', 'orders', 'subscriptions'],
@@ -498,51 +445,6 @@ def client(self,
             raise exceptions.ClientError("No such client.")
 
 
-def _get_filename_from_response(response) -> str:
-    """The name of the response resource.
-
-        The default is to use the content-disposition header value from the
-        response. If not found, falls back to resolving the name from the url
-        or generating a random name with the type from the response.
-        """
-    name = (_get_filename_from_headers(response.headers)
-            or _get_filename_from_url(response.url)
-            or _get_random_filename(response.headers.get('content-type')))
-    return name
-
-
-def _get_filename_from_headers(headers):
-    """Get a filename from the Content-Disposition header, if available.
-
-    :param headers dict: a ``dict`` of response headers
-    :returns: a filename (i.e. ``basename``)
-    :rtype: str or None
-    """
-    cd = headers.get('content-disposition', '')
-    match = re.search('filename="?([^"]+)"?', cd)
-    return match.group(1) if match else None
-
-
-def _get_filename_from_url(url: str) -> Optional[str]:
-    """Get a filename from a url.
-
-    Getting a name for Landsat imagery uses this function.
-    """
-    path = urlparse(url).path
-    name = path[path.rfind('/') + 1:]
-    return name or None
-
-
-def _get_random_filename(content_type=None) -> str:
-    """Get a pseudo-random, Planet-looking filename.
-    """
-    extension = mimetypes.guess_extension(content_type or '') or ''
-    characters = string.ascii_letters + '0123456789'
-    letters = ''.join(random.sample(characters, 8))
-    name = 'planet-{}{}'.format(letters, extension)
-    return name
-
-
 class AuthSession(BaseSession):
     """Synchronous connection to the Planet Auth service."""
 

planet/models.py

@@ -14,7 +14,9 @@
 # limitations under the License.
 """Manage data for requests and responses."""
 import logging
-from typing import AsyncGenerator, Callable, List
+import re
+from typing import AsyncGenerator, Callable, List, Optional
+from urllib.parse import urlparse
 
 import httpx
 
@@ -42,11 +44,67 @@ def status_code(self) -> int:
         """HTTP status code"""
         return self._http_response.status_code
 
+    @property
+    def filename(self) -> Optional[str]:
+        """Name of the download file.
+
+        The filename is None if the response does not represent a download.
+        """
+        filename = None
+
+        if self.length is not None:  # is a download file
+            filename = _get_filename_from_response(self._http_response)
+
+        return filename
+
+    @property
+    def length(self) -> Optional[int]:
+        """Length of the download file.
+
+        The length is None if the response does not represent a download.
+        """
+        LOGGER.warning('here')
+        try:
+            length = int(self._http_response.headers["Content-Length"])
+        except KeyError:
+            length = None
+        LOGGER.warning(length)
+        return length
+
     def json(self) -> dict:
         """Response json"""
         return self._http_response.json()
 
 
+def _get_filename_from_response(response) -> Optional[str]:
+    """The name of the response resource.
+
+        The default is to use the content-disposition header value from the
+        response. If not found, falls back to resolving the name from the url
+        or generating a random name with the type from the response.
+        """
+    name = (_get_filename_from_headers(response.headers)
+            or _get_filename_from_url(str(response.url)))
+    return name
+
+
+def _get_filename_from_headers(headers: httpx.Headers) -> Optional[str]:
+    """Get a filename from the Content-Disposition header, if available."""
+    cd = headers.get('content-disposition', '')
+    match = re.search('filename="?([^"]+)"?', cd)
+    return match.group(1) if match else None
+
+
+def _get_filename_from_url(url: str) -> Optional[str]:
+    """Get a filename from the  url.
+
+    Getting a name for Landsat imagery uses this function.
+    """
+    path = urlparse(url).path
+    name = path[path.rfind('/') + 1:]
+    return name or None
+
+
 class Paged:
     """Asynchronous iterator over results in a paged resource.
 

tests/integration/test_data_api.py

@@ -837,11 +837,15 @@ async def _stream_img():
     # populate request parameter to avoid respx cloning, which throws
     # an error caused by respx and not this code
     # https://github.com/lundberg/respx/issues/130
-    mock_resp = httpx.Response(HTTPStatus.OK,
-                               stream=_stream_img(),
-                               headers=img_headers,
-                               request='donotcloneme')
-    respx.get(dl_url).return_value = mock_resp
+    respx.get(dl_url).side_effect = [
+        httpx.Response(HTTPStatus.OK,
+                       headers=img_headers,
+                       request='donotcloneme'),
+        httpx.Response(HTTPStatus.OK,
+                       stream=_stream_img(),
+                       headers=img_headers,
+                       request='donotcloneme')
+    ]
 
     basic_udm2_asset = {
         "_links": {
@@ -863,7 +867,8 @@ async def _stream_img():
 
     path = await cl.download_asset(basic_udm2_asset,
                                    directory=tmpdir,
-                                   overwrite=overwrite)
+                                   overwrite=overwrite,
+                                   progress_bar=False)
     assert path.name == 'img.tif'
     assert path.is_file()