webrequests: Compress cache files

Author: mbriand

swattool/webrequests.py

@@ -2,12 +2,13 @@
 
 """Wrapper for requests module with cookies persistence and basic cache."""
 
+import gzip
 import hashlib
 import logging
 import pathlib
 import pickle
 import time
-from typing import Any
+from typing import Any, Optional
 
 import requests
 
@@ -61,6 +62,7 @@ def _get_cache_file_candidates(self, url: str) -> list[pathlib.Path]:
             filestem = hashname.hexdigest()
 
         candidates = [
+            utils.CACHEDIR / f"{filestem}.gz",
             utils.CACHEDIR / filestem,
 
             # For compatibility with old cache files
@@ -69,6 +71,32 @@ def _get_cache_file_candidates(self, url: str) -> list[pathlib.Path]:
 
         return candidates
 
+    def _try_load_cache(self, cachefile: pathlib.Path, max_cache_age: int
+                        ) -> Optional[str]:
+        if max_cache_age < 0:
+            use_cache = True
+        else:
+            age = time.time() - cachefile.stat().st_mtime
+            use_cache = age < max_cache_age
+
+        if use_cache:
+            if cachefile.suffix == ".gz":
+                with gzip.open(cachefile, mode='r') as gzfile:
+                    return gzfile.read(-1).decode()
+            else:
+                with cachefile.open('r') as file:
+                    return file.read(-1)
+
+        return None
+
+    def _create_cache_file(self, cachefile: pathlib.Path, data: str):
+        if cachefile.suffix == ".gz":
+            with gzip.open(cachefile, mode='w') as gzfile:
+                gzfile.write(data.encode())
+        else:
+            with cachefile.open('w') as file:
+                file.write(data)
+
     def get(self, url: str, max_cache_age: int = -1) -> str:
         """Do a GET request."""
         cache_candidates = self._get_cache_file_candidates(url)
@@ -77,24 +105,18 @@ def get(self, url: str, max_cache_age: int = -1) -> str:
 
         cache_olds = [file for file in cache_candidates if file.exists()]
         for cachefile in cache_olds:
-            if max_cache_age < 0:
-                use_cache = True
-            else:
-                age = time.time() - cachefile.stat().st_mtime
-                use_cache = age < max_cache_age
+            data = self._try_load_cache(cachefile, max_cache_age)
+            if data:
+                logger.debug("Loaded cache file for %s: %s", url, cachefile)
+                return data
 
-            if use_cache:
-                logger.debug("Loading cache file for %s: %s", url, cachefile)
-                with cachefile.open('r') as file:
-                    return file.read(-1)
-            else:
-                cachefile.unlink()
+            cachefile.unlink()
 
         logger.debug("Fetching %s, cache file will be %s", url, cache_new_file)
         req = self.session.get(url)
         req.raise_for_status()
-        with cache_new_file.open('w') as file:
-            file.write(req.text)
+
+        self._create_cache_file(cache_new_file, req.text)
 
         return req.text