Specify a custom user-agent string for all default requests.
[automanga.git] / manga / htcache.py
1 import os, hashlib, urllib.request, time
2 pj = os.path.join
3
4 class cache(object):
5     def __init__(self, dir):
6         self.dir = dir
7
8     def mangle(self, url):
9         n = hashlib.md5()
10         n.update(url.encode("ascii"))
11         return n.hexdigest()
12
13     def miss(self, url):
14         req = urllib.request.Request(url, headers={"User-Agent": "automanga/1"})
15         with urllib.request.urlopen(req) as s:
16             if s.headers.get("content-encoding") == "gzip":
17                 import gzip, io
18                 return gzip.GzipFile(fileobj=io.BytesIO(s.read()), mode="r").read()
19             return s.read()
20
21     def fetch(self, url, expire=3600):
22         path = pj(self.dir, self.mangle(url))
23         if os.path.exists(path):
24             if time.time() - os.stat(path).st_mtime < expire:
25                 with open(path, "rb") as f:
26                     return f.read()
27         data = self.miss(url)
28         if not os.path.isdir(self.dir):
29             os.makedirs(self.dir)
30         with open(path, "wb") as f:
31             f.write(data)
32         return data
33
34 home = os.getenv("HOME")
35 if home is None or not os.path.isdir(home):
36     raise Exception("Could not find home directory for HTTP caching")
37 default = cache(pj(home, ".manga", "htcache"))
38
39 def fetch(url, expire=3600):
40     return default.fetch(url, expire)