Merge pull request #664 from Soulsuck24/imgur

Serene-Arc · Sep 20, 2022 · 398f7b2 · 398f7b2
2 parents f4598c4 + cd05bc3
commit 398f7b2
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/bdfr/site_downloaders/download_factory.py b/bdfr/site_downloaders/download_factory.py
@@ -25,7 +25,7 @@ class DownloadFactory:
     @staticmethod
     def pull_lever(url: str) -> Type[BaseDownloader]:
         sanitised_url = DownloadFactory.sanitise_url(url)
-        if re.match(r'(i\.)?imgur.*\.gif.+$', sanitised_url):
+        if re.match(r'(i\.|m\.)?imgur', sanitised_url):
             return Imgur
         elif re.match(r'(i\.)?(redgifs|gifdeliverynetwork)', sanitised_url):
             return Redgifs
@@ -40,8 +40,6 @@ def pull_lever(url: str) -> Type[BaseDownloader]:
             return Gallery
         elif re.match(r'gfycat\.', sanitised_url):
             return Gfycat
-        elif re.match(r'(m\.)?imgur.*', sanitised_url):
-            return Imgur
         elif re.match(r'reddit\.com/r/', sanitised_url):
             return SelfPost
         elif re.match(r'(m\.)?youtu\.?be', sanitised_url):

diff --git a/bdfr/site_downloaders/imgur.py b/bdfr/site_downloaders/imgur.py
@@ -41,10 +41,12 @@ def _compute_image_url(self, image: dict) -> Resource:
 
     @staticmethod
     def _get_data(link: str) -> dict:
-        link = link.rstrip('?')
-        if re.match(r'(?i).*\.gif.+$', link):
-            link = link.replace('i.imgur', 'imgur')
-            link = re.sub('(?i)\\.gif.+$', '', link)
+        try:
+            imgur_id = re.match(r'.*/(.*?)(\..{0,})?$', link).group(1)
+            gallery = 'a/' if re.search(r'.*/(.*?)(gallery/|a/)', link) else ''
+            link = f'https://imgur.com/{gallery}{imgur_id}'
+        except AttributeError:
+            raise SiteDownloaderError(f'Could not extract Imgur ID from {link}')
 
         res = Imgur.retrieve_url(link, cookies={'over18': '1', 'postpagebeta': '0'})
 

diff --git a/tests/site_downloaders/test_download_factory.py b/tests/site_downloaders/test_download_factory.py
@@ -24,12 +24,12 @@
 @pytest.mark.parametrize(('test_submission_url', 'expected_class'), (
     ('https://www.reddit.com/r/TwoXChromosomes/comments/lu29zn/i_refuse_to_live_my_life'
      '_in_anything_but_comfort/', SelfPost),
-    ('https://i.imgur.com/bZx1SJQ.jpg', Direct),
+    ('https://i.imgur.com/bZx1SJQ.jpg', Imgur),
     ('https://i.redd.it/affyv0axd5k61.png', Direct),
-    ('https://imgur.com/3ls94yv.jpeg', Direct),
+    ('https://imgur.com/3ls94yv.jpeg', Imgur),
     ('https://i.imgur.com/BuzvZwb.gifv', Imgur),
     ('https://imgur.com/BuzvZwb.gifv', Imgur),
-    ('https://i.imgur.com/6fNdLst.gif', Direct),
+    ('https://i.imgur.com/6fNdLst.gif', Imgur),
     ('https://imgur.com/a/MkxAzeg', Imgur),
     ('https://i.imgur.com/OGeVuAe.giff', Imgur),
     ('https://www.reddit.com/gallery/lu93m7', Gallery),
@@ -40,7 +40,7 @@
     ('https://www.gifdeliverynetwork.com/repulsivefinishedandalusianhorse', Redgifs),
     ('https://youtu.be/DevfjHOhuFc', Youtube),
     ('https://m.youtube.com/watch?v=kr-FeojxzUM', Youtube),
-    ('https://i.imgur.com/3SKrQfK.jpg?1', Direct),
+    ('https://i.imgur.com/3SKrQfK.jpg?1', Imgur),
     ('https://dynasty-scans.com/system/images_images/000/017/819/original/80215103_p0.png?1612232781', Direct),
     ('https://m.imgur.com/a/py3RW0j', Imgur),
     ('https://v.redd.it/9z1dnk3xr5k61', VReddit),