From 5f77c4f091fa56b0f031f410c35510a016b66e9b Mon Sep 17 00:00:00 2001
From: Jason Dilworth <hello@jasondilworth.co.uk>
Date: Wed, 24 Jul 2024 23:07:30 +0100
Subject: [PATCH] Allow input of an author ID to download specific files.

---
 .gitignore              |  5 ++++-
 pyheadspace/__main__.py | 30 ++++++++++++++++++++++++++----
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 0fd5675..60fe462 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,4 +13,7 @@
 
 /pyheadspace/test.html
 /venv/
-.idea
\ No newline at end of file
+.idea
+
+.env
+launch.json
\ No newline at end of file
diff --git a/pyheadspace/__main__.py b/pyheadspace/__main__.py
index af5cccb..627a846 100644
--- a/pyheadspace/__main__.py
+++ b/pyheadspace/__main__.py
@@ -184,6 +184,7 @@ def get_pack_attributes(
     no_techniques: bool,
     no_meditation: bool,
     all_: bool = False,
+    author: Optional[int] = None,
 ):
     response = request_url(PACK_URL, id=pack_id)
     attributes: dict = response["data"]["attributes"]
@@ -209,11 +210,13 @@ def get_pack_attributes(
         if item["type"] == "orderedActivities":
             if not no_meditation:
                 id = item["relationships"]["activity"]["data"]["id"]
-                download_pack_session(id, duration, _pack_name, out=out)
+                download_pack_session(id, duration, _pack_name, out=out, author=author)
         elif item["type"] == "orderedTechniques":
             if not no_techniques:
                 id = item["relationships"]["technique"]["data"]["id"]
-                download_pack_techniques(id, pack_name=_pack_name, out=out)
+                download_pack_techniques(
+                    id, pack_name=_pack_name, out=out, author=author
+                )
 
 
 def get_signed_url(response: dict, duration: List[int]) -> dict:
@@ -265,8 +268,10 @@ def download_pack_session(
     pack_name: Optional[str],
     out: str,
     filename_suffix=None,
+    author: Optional[int] = None,
 ):
-    response = request_url(AUDIO_URL, id=id)
+    params = dict(authorId=author) if author else dict()
+    response = request_url(AUDIO_URL, id=id, params=params)
 
     signed_url = get_signed_url(response, duration=duration)
     for name, direct_url in signed_url.items():
@@ -281,8 +286,10 @@ def download_pack_techniques(
     pack_name: Optional[str] = None,
     out: str,
     filename_suffix=None,
+    author: Optional[int] = None,
 ):
-    response = request_url(TECHNIQUE_URL, id=technique_id)
+    params = dict(authorId=author) if author else dict()
+    response = request_url(TECHNIQUE_URL, id=technique_id, params=params)
     name = response["data"]["attributes"]["name"]
     if filename_suffix:
         name += filename_suffix
@@ -456,6 +463,16 @@ def get_legacy_id(new_id):
         " links of packs to exclude downloading. Every link should be on separate line."
     ),
 )
+@click.option(
+    "--author",
+    "-a",
+    type=int,
+    default=0,
+    help=(
+        "Use to choose the author/narrator that you'd like to download the files of."
+        "NOTE: If the author ID is not found, the default will download."
+    ),
+)
 @shared_cmd(COMMON_CMD)
 @shared_cmd(URL_GROUP_CMD)
 def pack(
@@ -467,6 +484,7 @@ def pack(
     url: str,
     all_: bool,
     exclude: str,
+    author: int,
 ):
     """
     Download headspace packs with techniques videos.
@@ -489,6 +507,7 @@ def pack(
             out=out,
             no_meditation=no_meditation,
             no_techniques=no_techniques,
+            author=author,
         )
     else:
         excluded = []
@@ -638,3 +657,6 @@ def login():
 
 
 session.close()
+
+if __name__ == "__main__":
+    cli()