From a937b72034d663f4ee7ddeb2898dbc9fa450be91 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mike=20F=C3=A4hrmann?= <mike_faehrmann@web.de>
Date: Sat, 28 Sep 2024 20:25:53 +0200
Subject: [PATCH] [ao3] add 'subscriptions' extractor (#6247)

---
 docs/supportedsites.md      |  2 +-
 gallery_dl/extractor/ao3.py | 45 ++++++++++++++++++++++++-------------
 test/results/ao3.py         | 34 +++++++++++++++-------------
 3 files changed, 50 insertions(+), 31 deletions(-)

diff --git a/docs/supportedsites.md b/docs/supportedsites.md
index e08f43a9..74112b3d 100644
--- a/docs/supportedsites.md
+++ b/docs/supportedsites.md
@@ -106,7 +106,7 @@ Consider all listed sites to potentially be NSFW.
 <tr>
     <td>Archive of Our Own</td>
     <td>https://archiveofourown.org/</td>
-    <td>Search Results, Series, Tag Searches, User Profiles, Bookmarks, Works</td>
+    <td>Search Results, Series, Subscriptions, Tag Searches, User Profiles, Bookmarks, Works</td>
     <td>Supported</td>
 </tr>
 <tr>
diff --git a/gallery_dl/extractor/ao3.py b/gallery_dl/extractor/ao3.py
index 1f570e83..d3ab8468 100644
--- a/gallery_dl/extractor/ao3.py
+++ b/gallery_dl/extractor/ao3.py
@@ -29,11 +29,31 @@ class Ao3Extractor(Extractor):
         self.login()
 
         base = self.root + "/works/"
-        data = {"_extractor": Ao3WorkExtractor}
+        data = {"_extractor": Ao3WorkExtractor, "type": "work"}
 
         for work_id in self.works():
             yield Message.Queue, base + work_id, data
 
+    def items_list(self, type, needle, part=True):
+        self.login()
+
+        base = self.root + "/"
+        data_work = {"_extractor": Ao3WorkExtractor, "type": "work"}
+        data_series = {"_extractor": Ao3SeriesExtractor, "type": "series"}
+        data_user = {"_extractor": Ao3UserExtractor, "type": "user"}
+
+        for item in self._pagination(self.groups[0], needle):
+            path = item.rpartition("/")[0] if part else item
+            url = base + path
+            if item.startswith("works/"):
+                yield Message.Queue, url, data_work
+            elif item.startswith("series/"):
+                yield Message.Queue, url, data_series
+            elif item.startswith("users/"):
+                yield Message.Queue, url, data_user
+            else:
+                self.log.warning("Unsupported %s type '%s'", type, path)
+
     def works(self):
         return self._pagination(self.groups[0])
 
@@ -284,19 +304,14 @@ class Ao3UserBookmarkExtractor(Ao3Extractor):
     example = "https://archiveofourown.org/users/USER/bookmarks"
 
     def items(self):
-        self.login()
+        return self.items_list("bookmark", '<span class="count"><a href="/')
 
-        base = self.root + "/"
-        data_work = {"_extractor": Ao3WorkExtractor}
-        data_series = {"_extractor": Ao3SeriesExtractor}
 
-        for item in self._pagination(
-                self.groups[0], '<span class="count"><a href="/'):
-            path = item.rpartition("/")[0]
-            url = base + path
-            if item.startswith("works/"):
-                yield Message.Queue, url, data_work
-            elif item.startswith("series/"):
-                yield Message.Queue, url, data_series
-            else:
-                self.log.warning("Unsupported bookmark type '%s'", path)
+class Ao3SubscriptionsExtractor(Ao3Extractor):
+    """Extractor for your AO3 account's subscriptions"""
+    subcategory = "subscriptions"
+    pattern = BASE_PATTERN + r"(/users/([^/?#]+)/subscriptions(?:/?\?.+)?)"
+    example = "https://archiveofourown.org/users/USER/subscriptions"
+
+    def items(self):
+        return self.items_list("subscription", '<dt>\n<a href="/', False)
diff --git a/test/results/ao3.py b/test/results/ao3.py
index d539bda4..4b363673 100644
--- a/test/results/ao3.py
+++ b/test/results/ao3.py
@@ -10,7 +10,6 @@ from gallery_dl.extractor import ao3
 __tests__ = (
 {
     "#url"     : "https://archiveofourown.org/works/47802076",
-    "#category": ("", "ao3", "work"),
     "#class"   : ao3.Ao3WorkExtractor,
     "#urls"    : "https://archiveofourown.org/downloads/47802076/The_Wildcard.pdf?updated_at=1720398424",
 
@@ -126,7 +125,6 @@ __tests__ = (
 
 {
     "#url"     : "https://archiveofourown.org/works/47802076",
-    "#category": ("", "ao3", "work"),
     "#class"   : ao3.Ao3WorkExtractor,
     "#options" : {"formats": ["epub", "mobi", "azw3", "pdf", "html"]},
     "#urls"    : (
@@ -138,9 +136,16 @@ __tests__ = (
     ),
 },
 
+{
+    "#url"     : "https://archiveofourown.org/works/12345",
+    "#comment" : "restricted work / login required",
+    "#class"   : ao3.Ao3WorkExtractor,
+    "#auth"    : True,
+    "#urls"    : "https://archiveofourown.org/downloads/12345/Unquenchable.pdf?updated_at=1716029699",
+},
+
 {
     "#url"     : "https://archiveofourown.org/series/1903930",
-    "#category": ("", "ao3", "series"),
     "#class"   : ao3.Ao3SeriesExtractor,
     "#urls"    : (
         "https://archiveofourown.org/works/26131546",
@@ -151,7 +156,6 @@ __tests__ = (
 
 {
     "#url"     : "https://archiveofourown.org/tags/Sunshine%20(Ghost%20Sweden%20Band)/works",
-    "#category": ("", "ao3", "tag"),
     "#class"   : ao3.Ao3TagExtractor,
     "#pattern" : ao3.Ao3WorkExtractor.pattern,
     "#range"   : "1-50",
@@ -160,7 +164,6 @@ __tests__ = (
 
 {
     "#url"     : "https://archiveofourown.org/works/search?work_search%5Bquery%5D=air+fire+ice+water",
-    "#category": ("", "ao3", "search"),
     "#class"   : ao3.Ao3SearchExtractor,
     "#pattern" : ao3.Ao3WorkExtractor.pattern,
     "#range"   : "1-50",
@@ -169,7 +172,6 @@ __tests__ = (
 
 {
     "#url"     : "https://archiveofourown.org/users/Fyrelass",
-    "#category": ("", "ao3", "user"),
     "#class"   : ao3.Ao3UserExtractor,
     "#urls"    : (
         "https://archiveofourown.org/users/Fyrelass/works",
@@ -178,39 +180,34 @@ __tests__ = (
 },
 {
     "#url"     : "https://archiveofourown.com/users/Fyrelass",
-    "#category": ("", "ao3", "user"),
     "#class"   : ao3.Ao3UserExtractor,
 },
 {
     "#url"     : "https://archiveofourown.net/users/Fyrelass",
-    "#category": ("", "ao3", "user"),
     "#class"   : ao3.Ao3UserExtractor,
 },
 {
     "#url"     : "https://ao3.org/users/Fyrelass",
-    "#category": ("", "ao3", "user"),
     "#class"   : ao3.Ao3UserExtractor,
 },
 
 {
     "#url"     : "https://archiveofourown.org/users/Fyrelass/profile",
-    "#category": ("", "ao3", "user"),
     "#class"   : ao3.Ao3UserExtractor,
 },
 
 {
     "#url"     : "https://archiveofourown.org/users/Fyrelass/pseuds/Aileen%20Autarkeia",
-    "#category": ("", "ao3", "user"),
     "#class"   : ao3.Ao3UserExtractor,
 },
 
 {
     "#url"     : "https://archiveofourown.org/users/Fyrelass/works",
-    "#category": ("", "ao3", "user-works"),
     "#class"   : ao3.Ao3UserWorksExtractor,
+    "#auth"    : False,
     "#urls"    : (
-        "https://archiveofourown.org/works/58979287",
         "https://archiveofourown.org/works/55035061",
+        "https://archiveofourown.org/works/58979287",
         "https://archiveofourown.org/works/52704457",
         "https://archiveofourown.org/works/52502743",
         "https://archiveofourown.org/works/52170409",
@@ -225,8 +222,8 @@ __tests__ = (
 
 {
     "#url"     : "https://archiveofourown.org/users/Fyrelass/series",
-    "#category": ("", "ao3", "user-series"),
     "#class"   : ao3.Ao3UserSeriesExtractor,
+    "#auth"    : False,
     "#urls"    : (
         "https://archiveofourown.org/series/3821575",
     ),
@@ -234,11 +231,18 @@ __tests__ = (
 
 {
     "#url"     : "https://archiveofourown.org/users/Fyrelass/bookmarks",
-    "#category": ("", "ao3", "user-bookmark"),
     "#class"   : ao3.Ao3UserBookmarkExtractor,
     "#pattern" : r"https://archiveofourown\.org/(work|serie)s/\d+",
     "#range"   : "1-50",
     "#count"   : 50,
 },
 
+{
+    "#url"     : "https://archiveofourown.org/users/mikf/subscriptions",
+    "#class"   : ao3.Ao3SubscriptionsExtractor,
+    "#auth"    : True,
+    "#pattern" : r"https://archiveofourown\.org/(work|serie|user)s/\w+",
+    "#count"   : range(20, 30),
+},
+
 )