add retry and exception support

py12306/app/app.py

@@ -15,6 +15,8 @@ class Config:
     PROJECT_DIR = os.path.abspath(__file__ + '/../../../') + '/'
     CONFIG_FILE = PROJECT_DIR + 'config.toml'
 
+    REQUEST_TIME_OUT = 5
+
     # Config
     REDIS = {
         'host': '127.0.0.1',
@@ -25,7 +27,7 @@ class Config:
     }
 
     # Redis keys
-    REDIS_PREFIX_KEY_TASKS = 'tasks:'
+    REDIS_PREFIX_KEY_TASKS = APP_NAME + ':tasks:'
 
     # REDIS_KEY_USER_TASKS = 'user_jobs'
 

py12306/app/query.py

@@ -0,0 +1,40 @@
+import re
+
+from py12306.app.app import Logger
+from py12306.lib.api import API_QUERY_INIT_PAGE
+from py12306.lib.exceptions import RetryException
+from py12306.lib.func import retry
+from py12306.lib.helper import ShareInstance
+from py12306.lib.request import Request
+
+
+class Query:
+    @classmethod
+    def task_train_ticket(cls, task: dict):
+        QueryTicket().query_with_info(task)
+
+
+class QueryTicket:
+    """
+    车票查询
+    """
+    api_type: str = None
+
+    def __init__(self):
+        self.session = Request()
+
+    def query_with_info(self, info: dict):
+        pass
+
+    @retry()
+    def get_query_api_type(self) -> str:
+        if QueryTicket.api_type:
+            return QueryTicket.api_type
+        response = self.session.get(API_QUERY_INIT_PAGE)
+        if response.status_code == 200:
+            res = re.search(r'var CLeftTicketUrl = \'(.*)\';', response.text)
+            try:
+                QueryTicket.api_type = res.group(1)
+            except IndexError:
+                raise RetryException('获取车票查询地址失败')
+        return self.get_query_api_type()

py12306/app/task.py

@@ -5,9 +5,10 @@ from py12306.lib.redis_lib import Redis
 
 def get_routes() -> dict:
     from py12306.app.user import User
+    from py12306.app.query import Query
     return {
-        'user': User.task_user,
-        'query': User.task_user,
+        'user_login': User.task_user_login,
+        'train_ticket': Query.task_train_ticket,
     }
 
 

py12306/app/user.py

@@ -4,9 +4,8 @@ from py12306.lib.helper import ShareInstance
 class User(ShareInstance):
 
     @classmethod
-    def task_user(cls, task: dict):
-        print(111)
-        a = 1
+    def task_user_login(cls, task: dict):
         pass
 
+
     pass

py12306/lib/api.py

@@ -0,0 +1,9 @@
+HOST_API = 'kyfw.12306.cn'
+BASE_API = 'https://' + HOST_API
+
+# LEFT_TICKETS = {
+#     "url": BASE_URL_OF_12306 + "/otn/{type}?leftTicketDTO.train_date={left_date}&leftTicketDTO.from_station={left_station}&leftTicketDTO.to_station={arrive_station}&purpose_codes=ADULT",
+# }
+
+
+API_QUERY_INIT_PAGE = BASE_API + '/otn/leftTicket/init'

py12306/lib/exceptions.py

@@ -0,0 +1,6 @@
+class RetryException(Exception):
+    pass
+
+
+class MaxRetryException(Exception):
+    pass

py12306/lib/func.py

@@ -22,3 +22,53 @@ def new_thread_with_jobs(jobs, wait=True, daemon=True, args=(), kwargs={}):
     if wait:
         for thread in threads:
             thread.join()
+
+
+def expand_class(cls, key, value, keep_old=True):
+    """
+    Expand class method
+    :param cls:
+    :param key:
+    :param value:
+    :param keep_old:
+    :return:
+    """
+    from types import MethodType
+
+    if keep_old:
+        setattr(cls, 'old_' + key, getattr(cls, key))
+    setattr(cls, key, MethodType(value, cls))
+    return cls
+
+
+def retry(num: int = 3):
+    """
+    Retry a func
+    :param num:
+    :return:
+    """
+    from py12306.lib.exceptions import RetryException, MaxRetryException
+    retry_num_key = '_retry_num'
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            retry_num = num
+            if retry_num_key in kwargs:
+                retry_num = kwargs.get(retry_num_key)
+                kwargs.pop('_retry_num')
+            try:
+                res = func(*args, **kwargs)
+            except RetryException as err:
+                retry_num -= 1
+                from py12306.app.app import Logger
+                Logger.warning('重试 %s, 剩余次数 %s' % (func.__name__, retry_num))
+                if retry_num > 0:
+                    kwargs[retry_num_key] = retry_num
+                    return wrapper(*args, **kwargs)
+                raise MaxRetryException(*err.args) from None
+
+            return res
+
+        return wrapper
+
+    return decorator

py12306/lib/helper.py

@@ -8,3 +8,15 @@ class ShareInstance():
         return cls.__session
 
 
+# Expand dict
+class Dict(dict):
+    def get(self, key, default=None, sep='.'):
+        keys = key.split(sep)
+        for i, key in enumerate(keys):
+            try:
+                value = self[key]
+                if len(keys[i + 1:]) and isinstance(value, Dict):
+                    return value.get(sep.join(keys[i + 1:]), default=default, sep=sep)
+                return value
+            except KeyError:
+                return self.dict_to_dict(default)

py12306/lib/request.py

@@ -0,0 +1,68 @@
+import requests
+from requests.exceptions import *
+from requests_html import HTMLSession, HTMLResponse
+
+from py12306.lib.func import expand_class
+
+requests.packages.urllib3.disable_warnings()
+
+
+class Request(HTMLSession):
+    """
+    请求处理类
+    """
+
+    def save_to_file(self, url, path):
+        response = self.get(url, stream=True)
+        with open(path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=1024):
+                f.write(chunk)
+        return response
+
+    @staticmethod
+    def _handle_response(response, **kwargs) -> HTMLResponse:
+        """
+        扩充 response
+        :param response:
+        :param kwargs:
+        :return:
+        """
+        response = HTMLSession._handle_response(response, **kwargs)
+        expand_class(response, 'json', Request.json)
+        return response
+
+    def add_response_hook(self, hook):
+        hooks = self.hooks['response']
+        if not isinstance(hooks, list):
+            hooks = [hooks]
+        hooks.append(hook)
+        self.hooks['response'] = hooks
+        return self
+
+    def json(self, default={}):
+        """
+        重写 json 方法，拦截错误
+        :return:
+        """
+        from py12306.lib.helper import Dict
+        try:
+            result = self.old_json()
+            return Dict(result)
+        except:
+            return Dict(default)
+
+    def request(self, *args, **kwargs):
+        try:
+            if 'timeout' not in kwargs:
+                from py12306.app.app import Config
+                kwargs['timeout'] = Config.REQUEST_TIME_OUT
+            response = super().request(*args, **kwargs)
+            return response
+        except RequestException as e:
+            if e.response:
+                response = e.response
+            else:
+                response = HTMLResponse(HTMLSession)
+                # response.status_code = 500
+                expand_class(response, 'json', Request.json)
+            return response

tests/test_query.py

@@ -0,0 +1,12 @@
+from tests.helper import BaseTest
+from py12306.app.query import QueryTicket
+
+
+class TestQueryTicket(BaseTest):
+    task = {
+        'name': 'admin',
+    }
+
+    def test_get_query_api_type(self):
+        res = QueryTicket().get_query_api_type()
+        self.assertEqual('leftTicket/query', res)