Add filter for polarization and additional calibration/noise/rfi downloads

asfsmd.py

@@ -67,7 +67,7 @@ def open(self, session=None):
         return self
 
 
-def query(products, auth=None):
+def query(products):
     """Query the specified Sentinel-1 products."""
     if isinstance(products, str):
         products = [products]
@@ -77,14 +77,24 @@ def query(products, auth=None):
 
 
 def download_annotations_core(urls, outdir=".", auth=None,
-                              block_size=BLOCKSIZE):
-    """Download Sentinel-1 annotationd for the specified product urls."""
+                              block_size=BLOCKSIZE, pol=None,
+                              do_calibration=False, do_noise=False,
+                              do_rfi=False):
+    """Download Sentinel-1 annotation for the specified product urls."""
     outdir = pathlib.Path(outdir)
 
-    patterns = [
-        "S1*.SAFE/manifest.safe",
-        "S1*.SAFE/annotation/s1*.xml",
-    ]
+    pol_filter = pol or ''  # empty string matches all polarizations
+
+    patterns = {
+        "S1*.SAFE/manifest.safe": '',
+        "S1*.SAFE/annotation/s1*.xml": pol_filter
+    }
+    if do_calibration:
+        patterns["S1*.SAFE/annotation/calibration/calibration*.xml"] = pol_filter
+    if do_noise:
+        patterns["S1*.SAFE/annotation/calibration/noise*.xml"] = pol_filter
+    if do_rfi:
+        patterns["S1*.SAFE/annotation/rfi/rfi*.xml"] = pol_filter
 
     with requests.Session() as session:
         session.auth = auth
@@ -108,10 +118,11 @@ def download_annotations_core(urls, outdir=".", auth=None,
                 with zipfile.ZipFile(fd) as zf:
                     components = []
                     for info in zf.filelist:
-                        for pattern in patterns:
+                        for pattern, filter in patterns.items():
                             if fnmatch.fnmatch(info.filename, pattern):
-                                components.append(info)
-                                break
+                                if filter in info.filename:
+                                    components.append(info)
+                                    break
 
                     component_iter = tqdm.tqdm(
                         components, unit="files", leave=False
@@ -133,18 +144,22 @@ def download_annotations_core(urls, outdir=".", auth=None,
                         _log.debug(f"{info.filename} extracted")
 
 
-def download_annotations(products, outdir=".", auth=None):
-    """Download annotationd for the specified Sentinel-1 products."""
-    results = query(products, auth=auth)
-    if len(results) != len(products):
-        warnings.warn(
-            f"only {len(results)} of the {len(products)} requested products "
-            f"found on the remote server"
-        )
+def download_annotations(products, outdir=".", auth=None, pol=None, urls=None,
+                         do_calibration=False, do_noise=False, do_rfi=False):
+    """Download annotations for the specified Sentinel-1 products."""
+    if urls is None:
+        results = query(products)
+        if len(results) != len(products):
+            warnings.warn(
+                f"only {len(results)} of the {len(products)} requested products "
+                f"found on the remote server"
+            )
 
-    urls = [item.properties["url"] for item in results]
+        urls = [item.properties["url"] for item in results]
 
-    download_annotations_core(urls, outdir=outdir, auth=auth)
+    download_annotations_core(urls, outdir=outdir, auth=auth, pol=pol,
+                              do_calibration=do_calibration, do_noise=do_noise,
+                              do_rfi=do_rfi)
 
 
 def _get_auth(*, user=None, pwd=None, hostname="urs.earthdata.nasa.gov"):
@@ -252,11 +267,11 @@ def _get_parser(subparsers=None):
         "-f",
         "--file-list",
         action="store_true",
-        help="read the list of products form file. "
+        help="read the list of products from a file. "
         "The file can be a JSON file (with '.json' extension) or a text file."
         "The text file is expected to contain one product name per line."
         "The json file can contain a list of products or a dictionary "
-        "containint a list of products for each key."
+        "containing a list of products for each key."
         "In this case the key is used as sub-folder name to store the "
         "corresponding products."
         "Example: <OUTDIR>/<KEY>/<PRODUCT>",
@@ -290,16 +305,48 @@ def _get_parser(subparsers=None):
         default=BLOCKSIZE,
         help="httpio block size in bytes (default: %(default)d)",
     )
+    # Optional filters
+    parser.add_argument(
+        "--pol",
+        "--polarization",
+        choices=["vv", "vh"],
+        type=str.lower,
+        help="Choose only one polarization to download. "
+        "If not provided both polarizations are downloaded."
+    ) 
+
+    # Additional file downloads
+    parser.add_argument(
+        "--do-calibration",
+        action="store_true",
+        help="Download calibration files."
+    )
+    parser.add_argument(
+        "--do-noise",
+        action="store_true",
+        help="Download noise calibration files."
+    )
+    parser.add_argument(
+        "--do-rfi",
+        action="store_true",
+        help="Download RFI files."
+    )
+    # Alternate way to get urls
+    parser.add_argument(
+        "--urls",
+        nargs="*",
+        help="Directly pass the urls from a query to the ASF API."
+    )
 
     # Positional arguments
     parser.add_argument(
         "inputs",
-        nargs="+",
+        nargs="*",
         metavar="INPUT",
         help="Sentinel-1 product name(s). "
         "If the '-f' flag is set then the argument is interpreted as "
         "the filename containing the list of products. "
-        "See '--file--list' option desctiption for more details",
+        "See '--file--list' option description for more details",
     )
 
     if subparsers is None:
@@ -347,7 +394,9 @@ def main(*argv):
                     assert isinstance(new_product, dict)
                     products_tree.update(new_product)
         else:
-            products_tree[""].extend(args.inputs)
+            # Ignore if user passed files with .zip or .SAFE extensions
+            inputs = [p.replace(".zip", "").replace(".SAFE", "") for p in args.inputs]
+            products_tree[""].extend(inputs)
 
         auth = _get_auth(user=args.username, pwd=args.password)
 
@@ -356,7 +405,11 @@ def main(*argv):
         for folder, products in items:
             pbar.set_description(folder if folder else 'DOWNLOAD')
             outpath = outroot / folder
-            download_annotations(products, outdir=outpath, auth=auth)
+            download_annotations(
+                products, outdir=outpath, auth=auth, pol=args.pol, urls=args.urls,
+                do_calibration=args.do_calibration, do_noise=args.do_noise,
+                do_rfi=args.do_rfi
+            )
 
     except Exception as exc:
         _log.critical(