refactor: remove file_name from Config

gadomski · gadomski · commit 26bd6b83be21 · 2023-08-14T16:15:43.000-06:00
diff --git a/src/stac_asset/_cli.py b/src/stac_asset/_cli.py
@@ -163,7 +163,6 @@ async def download_async(
         alternate_assets=alternate_assets,
         include=include,
         exclude=exclude,
-        file_name=file_name,
         s3_requester_pays=s3_requester_pays,
         s3_retry_mode=s3_retry_mode,
         s3_max_attempts=s3_max_attempts,
@@ -200,6 +199,7 @@ async def download() -> Union[Item, ItemCollection]:
             return await functions.download_item(
                 item,
                 directory_str,
+                file_name=file_name,
                 config=config,
                 queue=queue,
             )
@@ -211,6 +211,7 @@ async def download() -> Union[Item, ItemCollection]:
             return await functions.download_item_collection(
                 item_collection,
                 directory_str,
+                file_name=file_name,
                 config=config,
                 queue=queue,
             )
diff --git a/src/stac_asset/_download.py b/src/stac_asset/_download.py
@@ -9,7 +9,8 @@
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Type, Union
 
-from pystac import Asset, Collection, Item, STACObject
+import pystac.utils
+from pystac import Asset, Collection, Item, STACError
 from yarl import URL
 
 from .client import Client
@@ -28,6 +29,15 @@
     AnyQueue = Queue
 
 
+class WrappedError:
+    download: Download
+    error: Exception
+
+    def __init__(self, download: Download, error: Exception) -> None:
+        self.download = download
+        self.error = error
+
+
 @dataclass
 class Download:
     owner: Union[Item, Collection]
@@ -45,13 +55,6 @@ async def download(
             )
         except Exception as error:
             return WrappedError(self, error)
-        if "alternate" not in self.asset.extra_fields:
-            if not has_alternate_assets_extension(self.owner):
-                self.owner.stac_extensions.append(
-                    "https://stac-extensions.github.io/alternate-assets/v1.1.0/schema.json"
-                )
-            self.asset.extra_fields["alternate"] = {}
-        self.asset.extra_fields["alternate"]["from"] = {"href": self.asset.href}
         self.asset.href = str(self.path)
         return self
 
@@ -67,24 +70,43 @@ def __init__(self, config: Config) -> None:
         self.downloads = list()
         self.clients = dict()
 
-    async def add(self, stac_object: Union[Item, Collection], root: Path) -> None:
-        file_names: Set[str] = set()
-
+    async def add(
+        self, stac_object: Union[Item, Collection], root: Path, file_name: Optional[str]
+    ) -> None:
+        links = list()
+        for link in stac_object.links:
+            absolute_href = link.get_absolute_href()
+            if absolute_href:
+                link.target = absolute_href
+                links.append(link)
+        stac_object.links = links
+        # Will fail if the stac object doesn't have a self href and there's
+        # relative asset hrefs
+        stac_object = make_asset_hrefs_absolute(stac_object)
+
+        if file_name:
+            item_path = Path(root) / file_name
+            stac_object.set_self_href(str(item_path))
+        else:
+            item_path = None
+            stac_object.set_self_href(item_path)
+
+        asset_file_names: Set[str] = set()
         for key, asset in (
             (k, a)
             for k, a in stac_object.assets.items()
             if (not self.config.include or k in self.config.include)
             and (not self.config.exclude or k not in self.config.exclude)
         ):
             if self.config.asset_file_name_strategy == FileNameStrategy.FILE_NAME:
-                file_name = os.path.basename(URL(asset.href).path)
+                asset_file_name = os.path.basename(URL(asset.href).path)
             elif self.config.asset_file_name_strategy == FileNameStrategy.KEY:
-                file_name = key + Path(asset.href).suffix
+                asset_file_name = key + Path(asset.href).suffix
 
-            if file_name in file_names:
-                raise AssetOverwriteError(list(file_names))
+            if asset_file_name in asset_file_names:
+                raise AssetOverwriteError(list(asset_file_names))
             else:
-                file_names.add(file_name)
+                asset_file_names.add(asset_file_name)
 
             client_class = guess_client_class(asset, self.config)
             if client_class in self.clients:
@@ -98,7 +120,7 @@ async def add(self, stac_object: Union[Item, Collection], root: Path) -> None:
                     owner=stac_object,
                     key=key,
                     asset=asset,
-                    path=root / file_name,
+                    path=root / asset_file_name,
                     client=client,
                 )
             )
@@ -116,6 +138,7 @@ async def download(self, queue: Optional[AnyQueue]) -> None:
                 )
             )
         results = await asyncio.gather(*tasks)
+
         exceptions = set()
         for result in results:
             if isinstance(result, WrappedError):
@@ -195,17 +218,19 @@ def guess_client_class_from_href(href: str) -> Type[Client]:
         raise ValueError(f"could not guess client class for href: {href}")
 
 
-def has_alternate_assets_extension(stac_object: STACObject) -> bool:
-    return any(
-        extension.startswith("https://stac-extensions.github.io/alternate-assets")
-        for extension in stac_object.stac_extensions
-    )
-
-
-class WrappedError:
-    download: Download
-    error: Exception
-
-    def __init__(self, download: Download, error: Exception) -> None:
-        self.download = download
-        self.error = error
+def make_asset_hrefs_absolute(
+    stac_object: Union[Item, Collection]
+) -> Union[Item, Collection]:
+    # Copied from
+    # https://github.com/stac-utils/pystac/blob/381cf89fc25c15142fb5a187d905e22681de42a2/pystac/item.py#L309C3-L319C1
+    # until a fix for https://github.com/stac-utils/pystac/issues/1199 is
+    # released.
+    self_href = stac_object.get_self_href()
+    for asset in stac_object.assets.values():
+        if not pystac.utils.is_absolute_href(asset.href):
+            if self_href is None:
+                raise STACError(
+                    "Cannot make asset HREFs absolute if no self_href is set."
+                )
+            asset.href = pystac.utils.make_absolute_href(asset.href, self_href)
+    return stac_object
diff --git a/src/stac_asset/config.py b/src/stac_asset/config.py
@@ -37,12 +37,6 @@ class Config:
     Mutually exclusive with ``exclude``.
     """
 
-    file_name: Optional[str] = None
-    """The file name of the output item.
-
-    If not provided, the output item will not be saved.
-    """
-
     make_directory: bool = True
     """Whether to create the output directory.
 
diff --git a/src/stac_asset/functions.py b/src/stac_asset/functions.py
@@ -5,7 +5,7 @@
     TYPE_CHECKING,
     Any,
     Optional,
-    TypeVar,
+    Union,
 )
 
 import pystac.utils
@@ -25,6 +25,7 @@
 async def download_item(
     item: Item,
     directory: PathLikeObject,
+    file_name: Optional[str] = None,
     config: Optional[Config] = None,
     queue: Optional[AnyQueue] = None,
 ) -> Item:
@@ -33,6 +34,8 @@ async def download_item(
     Args:
         item: The :py:class:`pystac.Item`.
         directory: The output directory that will hold the items and assets.
+        file_name: The name of the item file to save. If not provided, will not
+            be saved.
         config: The download configuration
         queue: An optional queue to use for progress reporting
 
@@ -42,14 +45,24 @@ async def download_item(
     Raises:
         ValueError: Raised if the item doesn't have any assets.
     """
-    return await _download_stac_object(
-        item, directory, config=config or Config(), queue=queue
-    )
+    async with Downloads(config or Config()) as downloads:
+        await downloads.add(item, Path(directory), file_name)
+        await downloads.download(queue)
+
+    self_href = item.get_self_href()
+    if self_href:
+        _make_asset_hrefs_relative(item)
+        d = item.to_dict(include_self_link=True, transform_hrefs=False)
+        with open(self_href, "w") as f:
+            json.dump(d, f)
+
+    return item
 
 
 async def download_collection(
     collection: Collection,
     directory: PathLikeObject,
+    file_name: Optional[str] = None,
     config: Optional[Config] = None,
     queue: Optional[AnyQueue] = None,
 ) -> Collection:
@@ -61,6 +74,8 @@ async def download_collection(
     Args:
         collection: A pystac collection
         directory: The destination directory
+        file_name: The name of the collection file to save. If not provided,
+            will not be saved.
         config: The download configuration
         queue: An optional queue to use for progress reporting
 
@@ -70,14 +85,24 @@ async def download_collection(
     Raises:
         CantIncludeAndExclude: Raised if both include and exclude are not None.
     """
-    return await _download_stac_object(
-        collection, directory, config or Config(), queue=queue
-    )
+    async with Downloads(config or Config()) as downloads:
+        await downloads.add(collection, Path(directory), file_name)
+        await downloads.download(queue)
+
+    self_href = collection.get_self_href()
+    if self_href:
+        _make_asset_hrefs_relative(collection)
+        d = collection.to_dict(include_self_link=True, transform_hrefs=False)
+        with open(self_href, "w") as f:
+            json.dump(d, f)
+
+    return collection
 
 
 async def download_item_collection(
     item_collection: ItemCollection,
     directory: PathLikeObject,
+    file_name: Optional[str] = None,
     config: Optional[Config] = None,
     queue: Optional[AnyQueue] = None,
 ) -> ItemCollection:
@@ -86,6 +111,8 @@ async def download_item_collection(
     Args:
         item_collection: The item collection to download
         directory: The destination directory
+        file_name: The name of the item collection file to save. If not
+            provided, will not be saved.
         config: The download configuration
         queue: An optional queue to use for progress reporting
 
@@ -95,16 +122,14 @@ async def download_item_collection(
     Raises:
         CantIncludeAndExclude: Raised if both include and exclude are not None.
     """
-    if config is None:
-        config = Config()
-    async with Downloads(config) as downloads:
+    async with Downloads(config or Config()) as downloads:
         for item in item_collection.items:
             item.set_self_href(None)
             root = Path(directory) / item.id
-            await downloads.add(item, root)
+            await downloads.add(item, root, None)
         await downloads.download(queue)
-    if config.file_name:
-        dest_href = Path(directory) / config.file_name
+    if file_name:
+        dest_href = Path(directory) / file_name
         for item in item_collection.items:
             for asset in item.assets.values():
                 asset.href = pystac.utils.make_relative_href(
@@ -115,48 +140,9 @@ async def download_item_collection(
     return item_collection
 
 
-_T = TypeVar("_T", Collection, Item)
-
-
-async def _download_stac_object(
-    stac_object: _T,
-    directory: PathLikeObject,
-    config: Config,
-    queue: Optional[AnyQueue],
-) -> _T:
-    links = list()
-    for link in stac_object.links:
-        absolute_href = link.get_absolute_href()
-        if absolute_href:
-            link.target = absolute_href
-            links.append(link)
-    stac_object.links = links
-    # Will fail if the stac object doesn't have a self href and there's
-    # relative asset hrefs
-    stac_object = _make_asset_hrefs_absolute(stac_object)
-
-    if config.file_name:
-        item_path = Path(directory) / config.file_name
-        stac_object.set_self_href(str(item_path))
-    else:
-        item_path = None
-        stac_object.set_self_href(item_path)
-
-    async with Downloads(config) as downloads:
-        await downloads.add(stac_object, Path(directory))
-        await downloads.download(queue)
-
-    self_href = stac_object.get_self_href()
-    if self_href:
-        _make_asset_hrefs_relative(stac_object)
-        d = stac_object.to_dict(include_self_link=True, transform_hrefs=False)
-        with open(self_href, "w") as f:
-            json.dump(d, f)
-
-    return stac_object
-
-
-def _make_asset_hrefs_relative(stac_object: _T) -> _T:
+def _make_asset_hrefs_relative(
+    stac_object: Union[Item, Collection]
+) -> Union[Item, Collection]:
     # Copied from
     # https://github.com/stac-utils/pystac/blob/381cf89fc25c15142fb5a187d905e22681de42a2/pystac/item.py#L284C5-L298C20
     # until a fix for https://github.com/stac-utils/pystac/issues/1199 is
@@ -170,19 +156,3 @@ def _make_asset_hrefs_relative(stac_object: _T) -> _T:
                 )
             asset.href = pystac.utils.make_relative_href(asset.href, self_href)
     return stac_object
-
-
-def _make_asset_hrefs_absolute(stac_object: _T) -> _T:
-    # Copied from
-    # https://github.com/stac-utils/pystac/blob/381cf89fc25c15142fb5a187d905e22681de42a2/pystac/item.py#L309C3-L319C1
-    # until a fix for https://github.com/stac-utils/pystac/issues/1199 is
-    # released.
-    self_href = stac_object.get_self_href()
-    for asset in stac_object.assets.values():
-        if not pystac.utils.is_absolute_href(asset.href):
-            if self_href is None:
-                raise STACError(
-                    "Cannot make asset HREFs absolute if no self_href is set."
-                )
-            asset.href = pystac.utils.make_absolute_href(asset.href, self_href)
-    return stac_object
diff --git a/tests/test_functions.py b/tests/test_functions.py
diff --git a/tests/test_s3_client.py b/tests/test_s3_client.py
