From 3ccb39891ce5d2593aee56d469db68498dc2ab69 Mon Sep 17 00:00:00 2001
From: Christopher Barber <christopher.barber@analog.com>
Date: Thu, 24 Jun 2021 14:54:38 -0400
Subject: [PATCH] Support multi-platform downloads (#32)

---
 conda_mirror/conda_mirror.py | 74 +++++++++++++++++++++++++++++-------
 1 file changed, 61 insertions(+), 13 deletions(-)

diff --git a/conda_mirror/conda_mirror.py b/conda_mirror/conda_mirror.py
index 50b1143..e1b57c9 100644
--- a/conda_mirror/conda_mirror.py
+++ b/conda_mirror/conda_mirror.py
@@ -14,6 +14,7 @@
 import time
 import random
 from pprint import pformat
+from typing import List, Union
 
 import requests
 import yaml
@@ -130,7 +131,8 @@ def _make_arg_parser():
         ),
     )
     ap.add_argument(
-        "--target-directory", help="The place where packages should be mirrored to",
+        "--target-directory",
+        help="The place where packages should be mirrored to",
     )
     ap.add_argument(
         "--temp-directory",
@@ -144,9 +146,12 @@ def _make_arg_parser():
     )
     ap.add_argument(
         "--platform",
+        action="append",
+        default=[],
         help=(
             "The OS platform(s) to mirror. one of: {'linux-64', 'linux-32',"
-            "'osx-64', 'win-32', 'win-64'}"
+            "'osx-64', 'win-32', 'win-64'}. "
+            "May be specified multiple times."
         ),
     )
     ap.add_argument(
@@ -160,7 +165,9 @@ def _make_arg_parser():
         default=0,
     )
     ap.add_argument(
-        "--config", action="store", help="Path to the yaml config file",
+        "--config",
+        action="store",
+        help="Path to the yaml config file",
     )
     ap.add_argument(
         "--pdb",
@@ -176,7 +183,10 @@ def _make_arg_parser():
         help="Num of threads for validation. 1: Serial mode. 0: All available.",
     )
     ap.add_argument(
-        "--version", action="store_true", help="Print version and quit", default=False,
+        "--version",
+        action="store_true",
+        help="Print version and quit",
+        default=False,
     )
     ap.add_argument(
         "--dry-run",
@@ -333,6 +343,7 @@ def pdb_hook(exctype, value, traceback):
         else:
             url = "{}:{}".format(scheme, url[0])
         proxies = {scheme: url}
+
     return {
         "upstream_channel": args.upstream_channel,
         "target_directory": args.target_directory,
@@ -351,8 +362,7 @@ def pdb_hook(exctype, value, traceback):
 
 
 def cli():
-    """Thin wrapper around parsing the cli args and calling main with them
-    """
+    """Thin wrapper around parsing the cli args and calling main with them"""
     main(**_parse_and_format_args())
 
 
@@ -679,7 +689,7 @@ def _validate_or_remove_package(args):
     else:
         # Windows does not handle multiprocessing logging well
         # TODO: Fix this properly with a logging Queue
-        sys.stdout.write("Info: "+log_msg)
+        sys.stdout.write("Info: " + log_msg)
     package_path = os.path.join(package_directory, package)
     return _validate(
         package_path, md5=package_metadata.get("md5"), size=package_metadata.get("size")
@@ -690,7 +700,7 @@ def main(
     upstream_channel,
     target_directory,
     temp_directory,
-    platform,
+    platform: Union[str, List[str]],
     blacklist=None,
     whitelist=None,
     num_threads=1,
@@ -717,8 +727,8 @@ def main(
         The path on disk to an existing and writable directory to temporarily
         store the packages before moving them to the target_directory to
         apply checks
-    platform : str
-        The platform that you wish to mirror for. Common options are
+    platform : Union[str,List[str]]
+        The platforms that you wish to mirror. Common options are
         'linux-64', 'osx-64', 'win-64' and 'win-32'. Any platform is valid as
         long as the url resolves.
     blacklist : iterable of tuples, optional
@@ -805,6 +815,46 @@ def main(
         "blacklisted": set(),
         "to-mirror": set(),
     }
+
+    _platforms = [platform] if isinstance(platform, str) else platform
+
+    for platform in _platforms:
+        _download_platform(
+            summary,
+            upstream_channel,
+            target_directory,
+            temp_directory,
+            platform,
+            blacklist=blacklist,
+            whitelist=whitelist,
+            num_threads=num_threads,
+            dry_run=dry_run,
+            no_validate_target=no_validate_target,
+            minimum_free_space=minimum_free_space,
+            proxies=proxies,
+            ssl_verify=ssl_verify,
+            max_retries=max_retries,
+        )
+
+    return summary
+
+
+def _download_platform(
+    summary,
+    upstream_channel,
+    target_directory,
+    temp_directory,
+    platform: str,
+    blacklist=None,
+    whitelist=None,
+    num_threads=1,
+    dry_run=False,
+    no_validate_target=False,
+    minimum_free_space=0,
+    proxies=None,
+    ssl_verify=None,
+    max_retries=100,
+):
     # Implementation:
     if not os.path.exists(os.path.join(target_directory, platform)):
         os.makedirs(os.path.join(target_directory, platform))
@@ -880,7 +930,7 @@ def main(
     summary["to-mirror"].update(to_mirror)
     if dry_run:
         logger.info("Dry run complete. Exiting")
-        return summary
+        return
 
     # 6. for each download:
     # a. download to temp file
@@ -975,8 +1025,6 @@ def main(
         noarch_repodata = {"info": {}, "packages": {}}
         _write_repodata(noarch_path, noarch_repodata)
 
-    return summary
-
 
 def _write_repodata(package_dir, repodata_dict):
     data = json.dumps(repodata_dict, indent=2, sort_keys=True)