operations/files.download: add extra optional arguments for wget or curl

Author: jgelens

pyinfra/operations/files.py

@@ -83,6 +83,8 @@ def download(
     insecure=False,
     proxy: str | None = None,
     temp_dir: str | Path | None = None,
+    extra_curl_args: dict[str, str] | None = None,
+    extra_wget_args: dict[str, str] | None = None,
 ):
     """
     Download files from remote locations using ``curl`` or ``wget``.
@@ -102,6 +104,8 @@ def download(
     + insecure: disable SSL verification for the HTTP request
     + proxy: simple HTTP proxy through which we can download files, form `http://<yourproxy>:<port>`
     + temp_dir: use this custom temporary directory during the download
+    + extra_curl_args: optional dictionary with custom arguments for curl
+    + extra_wget_args: optional dictionary with custom arguments for wget
 
     **Example:**
 
@@ -164,6 +168,14 @@ def download(
         curl_args: list[Union[str, StringCommand]] = ["-sSLf"]
         wget_args: list[Union[str, StringCommand]] = ["-q"]
 
+        if extra_curl_args:
+            for key, value in extra_curl_args.items():
+                curl_args.append(StringCommand(key, QuoteString(value)))
+
+        if extra_wget_args:
+            for key, value in extra_wget_args.items():
+                wget_args.append(StringCommand(key, QuoteString(value)))
+
         if proxy:
             curl_args.append(f"--proxy {proxy}")
             wget_args.append("-e use_proxy=yes")

tests/operations/files.download/download_extra_args_curl.json

@@ -0,0 +1,22 @@
+{
+    "args": ["http://myfile"],
+    "kwargs": {
+        "dest": "/home/myfile",
+        "extra_curl_args": {
+            "--user": "abc:def"
+        }
+    },
+    "facts": {
+        "server.Date": "datetime:2015-01-01T00:00:00",
+        "files.File": {
+            "path=/home/myfile": null
+        },
+        "server.Which": {
+            "command=curl": "yes"
+        }
+    },
+    "commands": [
+        "curl -sSLf --user abc:def http://myfile -o _tempfile_",
+        "mv _tempfile_ /home/myfile"
+    ]
+}

tests/operations/files.download/download_extra_args_wget.json

@@ -0,0 +1,23 @@
+{
+    "args": ["http://myfile"],
+    "kwargs": {
+        "dest": "/home/myfile",
+        "extra_wget_args": {
+            "--user": "abc:def"
+        }
+    },
+    "facts": {
+        "server.Date": "datetime:2015-01-01T00:00:00",
+        "files.File": {
+            "path=/home/myfile": null
+        },
+        "server.Which": {
+            "command=curl": null,
+            "command=wget": "yes"
+        }
+    },
+    "commands": [
+        "wget -q --user abc:def http://myfile -O _tempfile_ || ( rm -f _tempfile_ ; exit 1 )",
+        "mv _tempfile_ /home/myfile"
+    ]
+}