add Web helper & logger

hexlib/db.py

@@ -3,11 +3,11 @@ import sqlite3
 import traceback
 
 import psycopg2
-import redis
-import orjson as json
 import umsgpack
 from psycopg2.errorcodes import UNIQUE_VIOLATION
 
+from hexlib.env import get_redis
+
 
 class PersistentState:
     """Quick and dirty persistent dict-like SQLite wrapper"""
@@ -25,8 +25,8 @@ class PersistentState:
 class VolatileState:
     """Quick and dirty volatile dict-like redis wrapper"""
 
-    def __init__(self, prefix, **redis_args):
-        self.rdb = redis.Redis(**redis_args)
+    def __init__(self, prefix, redis_db):
+        self.rdb = redis_db
         self.prefix = prefix
 
     def __getitem__(self, table):
@@ -36,8 +36,10 @@ class VolatileState:
 class VolatileQueue:
     """Quick and dirty volatile queue-like redis wrapper"""
 
-    def __init__(self, key, **redis_args):
-        self.rdb = redis.Redis(**redis_args)
+    def __init__(self, key, redis_db=None):
+        if redis_db is None:
+            redis_db = get_redis()
+        self.rdb = redis_db
         self.key = key
 
     def put(self, item):
@@ -52,8 +54,8 @@ class VolatileQueue:
 class VolatileBooleanState:
     """Quick and dirty volatile dict-like redis wrapper for boolean values"""
 
-    def __init__(self, prefix, **redis_args):
-        self.rdb = redis.Redis(**redis_args)
+    def __init__(self, prefix, redis_db):
+        self.rdb = redis_db
         self.prefix = prefix
 
     def __getitem__(self, table):
@@ -184,7 +186,6 @@ class Table:
                 pass
 
 
-
 def _sqlite_type(value):
     if isinstance(value, int):
         return "integer"

hexlib/env.py

@@ -0,0 +1,22 @@
+import redis
+import os
+
+from hexlib.log import stdout_logger
+from hexlib.web import Web
+
+
+def get_redis():
+    return redis.Redis(
+        host=os.environ.get("REDIS_HOST", "localhost"),
+        port=int(os.environ.get("REDIS_PORT", 6379))
+    )
+
+
+def get_web():
+    return Web(
+        proxy=os.environ.get("PROXY", None),
+        rps=os.environ.get("RPS", 1),
+        logger=stdout_logger,
+        cookie_file=os.environ.get("COOKIE_FILE", None),
+        retry_codes=set(os.environ.get("RETRY_CODES", "").split(","))
+    )

hexlib/log.py

@@ -0,0 +1,51 @@
+import logging
+import os
+import sys
+from logging import StreamHandler
+
+DATE_FMT = "%Y-%m-%d %H:%M:%S"
+
+
+class ColorFormatter(logging.Formatter):
+
+    def __init__(self, fmt):
+        super().__init__()
+
+        grey = "\x1b[38;21m"
+        yellow = "\x1b[33;21m"
+        red = "\x1b[31;21m"
+        bold_red = "\x1b[31;1m"
+        reset = "\x1b[0m"
+
+        self.formats = {
+            logging.DEBUG: logging.Formatter(grey + fmt + reset, datefmt=DATE_FMT),
+            logging.INFO: logging.Formatter(grey + fmt + reset, datefmt=DATE_FMT),
+            logging.WARNING: logging.Formatter(yellow + fmt + reset, datefmt=DATE_FMT),
+            logging.ERROR: logging.Formatter(red + fmt + reset, datefmt=DATE_FMT),
+            logging.CRITICAL: logging.Formatter(bold_red + fmt + reset, datefmt=DATE_FMT)
+        }
+
+    def format(self, record):
+        return self.formats[record.levelno].format(record)
+
+
+stdout_logger = logging.getLogger("default")
+
+if os.environ.get("LOG_LEVEL", "debug") == "debug":
+    stdout_logger.setLevel(logging.DEBUG)
+
+for h in stdout_logger.handlers:
+    stdout_logger.removeHandler(h)
+
+handler = StreamHandler(sys.stdout)
+if os.environ.get("LOG_THREAD_NAME", "0") == "1":
+    fmt = "%(asctime)s %(levelname)-5s>%(threadName)s %(message)s"
+else:
+    fmt = "%(asctime)s %(levelname)-5s>%(message)s"
+
+if os.environ.get("LOG_COLORS", "1") == "1":
+    handler.formatter = ColorFormatter(fmt)
+else:
+    handler.formatter = logging.Formatter(fmt, datefmt='%Y-%m-%d %H:%M:%S')
+stdout_logger.addHandler(handler)
+logger = stdout_logger

hexlib/web.py

@@ -4,12 +4,16 @@ import os
 from datetime import datetime
 from base64 import b64encode, b64decode
 from http.cookiejar import Cookie
+from time import time
+
 import requests
 import orjson as json
 
 from dateutil.parser import parse
 from requests.cookies import RequestsCookieJar
 
+from hexlib.misc import rate_limit, retry
+
 
 def cookie_from_string(text: str, domain: str) -> Cookie:
     tokens = [t.strip() for t in text.split(";")]
@@ -104,3 +108,63 @@ def download_file(url, destination, session=None, headers=None, overwrite=False,
             if err_cb:
                 err_cb(e)
             retries -= 1
+
+class Web:
+    def __init__(self, proxy=None, rps=1, retries=3, logger=None, cookie_file=None, retry_codes=None, session=None):
+        self._cookie_file = cookie_file
+        self._proxy = proxy
+        self._logger = logger
+        self._current_req = None
+        if retry_codes is None:
+            retry_codes = {502, 504, 522, 524, 429}
+        self._retry_codes = retry_codes
+
+        if session is None:
+            session = requests.session()
+
+        self._session = session
+
+        if self._cookie_file:
+            self._session.cookies = load_cookiejar(cookie_file)
+
+        if self._proxy:
+            self._session.proxies = {
+                "http": proxy,
+                "https": proxy,
+            }
+
+        @rate_limit(rps)
+        @retry(retries, callback=self._error_callback)
+        def get(url, **kwargs):
+            self._current_req = "GET", url, kwargs
+            r = self._session.get(url, **kwargs)
+
+            if r.status_code in self._retry_codes:
+                raise Exception(f"HTTP {r.status_code}")
+            return r
+
+        self._get = get
+
+    def _error_callback(self, e):
+        self._logger.critical(f"{self._format_url(*self._current_req)}: {e}")
+
+    def _format_url(self, method, url, kwargs, r=None):
+        if "params" in kwargs and kwargs["params"]:
+            return "%s %s?%s <%s>" % (method, url, "&".join(f"{k}={v}" for k, v in kwargs["params"].items()),
+                                      r.status_code if r else "ERR")
+        else:
+            return "%s %s <%s>" % (method, url, r.status_code if r else "ERR",)
+
+    def get(self, url, **kwargs):
+
+        time_start = time()
+        r = self._get(url, **kwargs)
+
+        if self._cookie_file:
+            save_cookiejar(self._session.cookies, self._cookie_file)
+
+        if self._logger and r is not None:
+            self._logger.debug(self._format_url("GET", url, kwargs, r) + " %.2fs" % (time() - time_start))
+        return r
+
+

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup
 
 setup(
     name="hexlib",
-    version="1.28",
+    version="1.29",
     description="Misc utility methods",
     author="simon987",
     author_email="me@simon987.net",

test/test_VolatileState.py

@@ -1,4 +1,5 @@
 from unittest import TestCase
+
 from hexlib.db import VolatileState, VolatileBooleanState, VolatileQueue
 
 

test/test_download_file.py

@@ -1,5 +1,5 @@
-from unittest import TestCase
 import os
+from unittest import TestCase
 
 from hexlib.web import download_file
 

test/test_retry.py

@@ -1,5 +1,5 @@
 from unittest import TestCase
-from hexlib.db import VolatileState, VolatileBooleanState
+
 from hexlib.misc import retry