Repos / pytaku / 3b685d7f6f
commit 3b685d7f6f3197ab7c1d248ece298db6f1b61b14
Author: Bùi Thành Nhân <hi@imnhan.com>
Date:   Sun Aug 16 15:14:49 2020 +0700

    use per-site http session object

diff --git a/src/mangoapi/base_site.py b/src/mangoapi/base_site.py
index afa801d..5cbbc36 100644
--- a/src/mangoapi/base_site.py
+++ b/src/mangoapi/base_site.py
@@ -1,12 +1,20 @@
 import functools
 from abc import ABC, abstractmethod
 
+import requests
+
 
 class Site(ABC):
     def __init__(self):
-        self._cookies = None
         self.username = None
         self.password = None
+        self.is_logged_in = False
+        self.session = requests.Session()
+        self.session.headers.update(
+            {
+                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/78.0.3904.108 Safari/537.36"
+            }
+        )
 
     @abstractmethod
     def get_title(self, title_id):
@@ -45,10 +53,12 @@ def requires_login(func):
 
     @functools.wraps(func)
     def wrapper(self, *args, **kwargs):
-        if self._cookies is None:
+        # TODO: replace is_logged_in flag check with actual "if rejected then try
+        # logging in" logic, just in case login cookies expire.
+        if not self.is_logged_in:
             assert self.username
             assert self.password
-            self._cookies = self.login(self.username, self.password)
+            self.login(self.username, self.password)
         return func(self, *args, **kwargs)
 
     return wrapper
diff --git a/src/mangoapi/mangadex.py b/src/mangoapi/mangadex.py
index 4894744..90c404b 100644
--- a/src/mangoapi/mangadex.py
+++ b/src/mangoapi/mangadex.py
@@ -1,15 +1,13 @@
 import html
 import re
 
-import requests
-
 from mangoapi.base_site import Site, requires_login
 
 
 class Mangadex(Site):
     def get_title(self, title_id):
         url = f"https://mangadex.org/api/?id={title_id}&type=manga"
-        md_resp = requests.get(url)
+        md_resp = self.session.get(url)
         assert md_resp.status_code == 200, md_resp.text
         md_json = md_resp.json()
         assert md_json["status"] == "OK"
@@ -41,7 +39,7 @@ def get_title(self, title_id):
         return title
 
     def get_chapter(self, title_id, chapter_id):
-        md_resp = requests.get(
+        md_resp = self.session.get(
             f"https://mangadex.org/api/?id={chapter_id}&type=chapter&saver=0"
         )
         assert md_resp.status_code == 200, md_resp.text
@@ -65,9 +63,7 @@ def get_chapter(self, title_id, chapter_id):
 
     @requires_login
     def search_title(self, query):
-        md_resp = requests.get(
-            f"https://mangadex.org/quick_search/{query}", cookies=self._cookies,
-        )
+        md_resp = self.session.get(f"https://mangadex.org/quick_search/{query}")
         assert md_resp.status_code == 200, md_resp.text
 
         matches = TITLES_PATTERN.findall(md_resp.text)
@@ -83,22 +79,19 @@ def search_title(self, query):
         return titles
 
     def login(self, username, password):
-        """
-        Returns cookies of a logged in user.
-        """
         form_data = {
             "login_username": username,
             "login_password": password,
             "two_factor": "",
             "remember_me": "1",
         }
-        md_resp = requests.post(
+        md_resp = self.session.post(
             "https://mangadex.org/ajax/actions.ajax.php?function=login",
             data=form_data,
             headers={"X-Requested-With": "XMLHttpRequest"},
         )
         assert md_resp.status_code == 200, md_resp.text
-        return dict(md_resp.cookies)
+        self.is_logged_in = True
 
     def title_cover(self, title_id, cover_ext):
         return f"https://mangadex.org/images/manga/{title_id}.{cover_ext}"
diff --git a/src/mangoapi/mangasee.py b/src/mangoapi/mangasee.py
index 0df9a00..5f2080b 100644
--- a/src/mangoapi/mangasee.py
+++ b/src/mangoapi/mangasee.py
@@ -1,9 +1,8 @@
 import json
 import re
 
-import requests
-
 import apsw
+
 from mangoapi.base_site import Site
 
 regexes = {
@@ -22,10 +21,11 @@ class Mangasee(Site):
     search_table = None
 
     def __init__(self, keyval_store=None):
+        super().__init__()
         self.keyval_store = keyval_store
 
     def get_title(self, title_id):
-        resp = requests.get(f"https://mangasee123.com/manga/{title_id}", timeout=3)
+        resp = self.session.get(f"https://mangasee123.com/manga/{title_id}", timeout=3)
         assert resp.status_code == 200
         html = resp.text
         name = regexes["title_name"].search(html).group(1).strip()
@@ -54,7 +54,7 @@ def get_title(self, title_id):
         }
 
     def get_chapter(self, title_id, chapter_id):
-        resp = requests.get(
+        resp = self.session.get(
             f"https://mangasee123.com/read-online/{title_id}-chapter-{chapter_id}.html"
         )
         assert resp.status_code == 200
@@ -108,7 +108,7 @@ def search_title(self, query):
                 )
             if not titles:
                 print("Fetching mangasee title list...", end="")
-                resp = requests.get("https://mangasee123.com/_search.php")
+                resp = self.session.get("https://mangasee123.com/_search.php")
                 print(" done")
                 titles = resp.json()
                 self.keyval_store.set("mangasee_titles", resp.text)
diff --git a/tests/mangoapi/test_mangadex.py b/tests/mangoapi/test_mangadex.py
index 6be93fb..b75b5b6 100644
--- a/tests/mangoapi/test_mangadex.py
+++ b/tests/mangoapi/test_mangadex.py
@@ -1,7 +1,6 @@
-import os
+from pytaku.conf import config
 
 from mangoapi.mangadex import Mangadex
-from pytaku.conf import config
 
 
 def test_get_title():