Add parameters to fetch subdomains, full list

Refactor caching system
Require caching
Version bump

Author: VictorLeP

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="tranco",
-    version="0.7.1",
+    version="0.8.0",
     author="Victor Le Pochat",
     author_email="[email protected]",
     description="Tranco: A Research-Oriented Top Sites Ranking Hardened Against Manipulation",

tests/conftest.py

@@ -5,7 +5,6 @@
 @pytest.fixture(scope="session")
 def tranco():
     t = Tranco(cache=True, cache_dir='.tranco')
-    t.list()  # prefetch list
     return t
 
 

tests/test_download.py

@@ -0,0 +1,3 @@
+def test_daily(tranco):
+    l = tranco.list(date="2024-01-01")
+    assert l.list_id == "V929N"

tests/test_parameters.py

@@ -0,0 +1,15 @@
+def test_full(tranco):
+    l = tranco.list(list_id="93N82", full=True)
+    assert len(l.list) == 8469
+
+
+def test_subdomains(tranco):
+    l = tranco.list(date="2024-01-01", subdomains=True)
+    assert l.list_id == "G6Y6K"
+
+
+def test_top_1m_after_full(tranco):
+    lf = tranco.list(list_id="G6Y6K", full=True)
+    assert len(lf.list) > 1000000
+    ln = tranco.list(list_id="G6Y6K", full=False)
+    assert len(ln.list) == 1000000

tests/test_list.py tests/test_rank.pytests/test_list.py renamed to tests/test_rank.py

@@ -1,14 +1,5 @@
 import uuid
 
-
-def test_top_10(tranco_list):
-    assert len(tranco_list.top(10)) == 10
-
-
-def test_top_1000000(tranco_list):
-    assert len(tranco_list.top(1000000)) == 1000000
-
-
 def test_domain_rank(tranco_list):
     top_1 = tranco_list.top(1)[0]
     assert tranco_list.rank(top_1) == 1

tests/test_top.py

@@ -0,0 +1,8 @@
+def test_top_10(tranco_list):
+    assert len(tranco_list.top(10)) == 10
+
+
+def test_top_1000000(tranco_list):
+    assert len(tranco_list.top(1000000)) == 1000000
+
+

tranco/tranco.py

@@ -1,11 +1,16 @@
+import json
 from datetime import datetime, timedelta
 from io import BytesIO
+from itertools import islice
+
 import requests
 import os
 import platform
 import zipfile
+from warnings import warn
+from enum import IntEnum
 
-VERSION = '0.7.1'
+VERSION = '0.8.0'
 
 class TrancoList():
     def __init__(self, date, list_id, lst):
@@ -20,91 +25,156 @@ def top(self, num=1000000):
     def rank(self, domain):
         return self.list.get(domain, -1)
 
+class TrancoCacheType(IntEnum):
+    NOT_CACHED = 0
+    CACHED_NOT_FULL = 1
+    CACHED_FULL = 2
+
 class Tranco():
     def __init__(self, **kwargs):
         """
         :param kwargs:
-            cache: <bool> enables/disables caching, default: True
             cache_dir: <str> directory used to cache Tranco top lists, default: cwd + .tranco/
             account_email: <str> Account email address: retrieve from https://tranco-list.eu/account
             api_key: <str> API key: retrieve from https://tranco-list.eu/account
         """
 
-        self.should_cache = kwargs.get('cache', True)
+        # Caching is required.
         self.cache_dir = kwargs.get('cache_dir', None)
         if self.cache_dir is None:
             cwd = os.getcwd()
             self.cache_dir = os.path.join(cwd, '.tranco')
-
-        if self.should_cache and not os.path.exists(self.cache_dir):
+        if not os.path.exists(self.cache_dir):
             os.mkdir(self.cache_dir)
+        self.cache_metadata = {}
+        self._load_cache_metadata()
 
         self.account_email = kwargs.get('account_email')
         self.api_key = kwargs.get('api_key')
 
         self.session = requests.Session()
         self.session.headers.update({'User-Agent': 'Python/{} python-requests/{} tranco-python/{}'.format(platform.python_version(), requests.__version__, VERSION)})
 
-    def _cache_path(self, date):
-        return os.path.join(self.cache_dir, date + '-DEFAULT.csv')
+    def _cache_metadata_path(self):
+        return os.path.join(self.cache_dir, 'metadata.json')
+
+    def _cache_path(self, list_id):
+        return os.path.join(self.cache_dir, '{}.csv'.format(list_id))
+
+    def _load_cache_metadata(self):
+        if not os.path.exists(self._cache_metadata_path()):
+            self._write_cache_metadata()
+        with open(self._cache_metadata_path(), "rt") as f:
+            self.cache_metadata = json.load(f)
+
+    def _write_cache_metadata(self):
+        with open(self._cache_metadata_path(), 'wt') as f:
+            json.dump(self.cache_metadata, f)
+
+    def _get_list_cache(self, list_id):
+        return self.cache_metadata.get(list_id, TrancoCacheType.NOT_CACHED)
+
+    def _is_cached(self, list_id, full=False):
+        if not list_id:
+            raise ValueError("You must pass a list ID to cache a list.")
+        list_cache = self._get_list_cache(list_id)
+        if list_cache == TrancoCacheType.NOT_CACHED:
+            return False
+
+        if full and (list_cache == TrancoCacheType.CACHED_NOT_FULL):  # need full, but full not present
+            return False
+        return True
+
+    def _add_to_cache(self, list_id=None, full=False):
+        if not list_id:
+            raise ValueError("You must pass a list ID to cache a list.")
+        self.cache_metadata[list_id] = max(TrancoCacheType.CACHED_FULL if full else TrancoCacheType.CACHED_NOT_FULL, self._get_list_cache(list_id))
+        self._write_cache_metadata()
+
+    def clear_cache(self):
+        for f in os.listdir(self.cache_dir):
+            os.remove(os.path.join(self.cache_dir, f))
+        self._load_cache_metadata()
 
-    def list(self, date=None, list_id=None):
+    def list(self, date=None, list_id=None, subdomains=False, full=False):
+        """
+        Retrieve a Tranco top list.
+        :param date: Get the daily list for this date. If not given, the latest list is returned.
+                     Combine with `subdomains` to select whether subdomains are included.
+        :param list_id: Get the list with this ID. If neither the list ID nor date are given, the latest list is returned.
+        :param subdomains: Include subdomains in the list. Only relevant when requesting a daily list. Default: False.
+        :param full: Retrieve the full list (else only the top million). Default: False.
+        :return: TrancoList object for the requested list.
+        """
         if date and list_id:
             raise ValueError("You can't pass a date as well as a list ID.")
+        if list_id and subdomains:
+            warn("Subdomains parameter is ignored when passing a list ID.")
 
         if not list_id:
             if (not date) or (date == 'latest'):  # no arguments given: default to latest list
                 yesterday = (datetime.utcnow() - timedelta(days=1))
                 date = yesterday.strftime('%Y-%m-%d')
-            list_id = self._get_list_id_for_date(date)
+            list_id = self._get_list_id_for_date(date, subdomains=subdomains)
 
-        if self.should_cache and os.path.exists(self._cache_path(list_id)):
-            with open(self._cache_path(list_id)) as f:
-                top_list_text = f.read()
-        else:
-            top_list_text = self._download_zip_file(list_id)
+        if not self._is_cached(list_id, full):
+            self._download_file(list_id, full)  # download list and load into cache
+        with open(self._cache_path(list_id)) as f:  # read list from cache
+            if full:
+                top_list_lines = f.read().splitlines()
+            else:
+                top_list_lines = list(islice(f, 1000000))
 
-        return TrancoList(date, list_id, list(map(lambda x: x[x.index(',') + 1:], top_list_text.splitlines())))
+        return TrancoList(date, list_id, list(map(lambda x: x[x.index(',') + 1:], top_list_lines)))
 
-    def _get_list_id_for_date(self, date):
-        r1 = self.session.get('https://tranco-list.eu/daily_list_id?date={}'.format(date))
+    def _get_list_id_for_date(self, date, subdomains=False):
+        r1 = self.session.get('https://tranco-list.eu/daily_list_id?date={}&subdomains={}'.format(date, str(subdomains).lower()))
         if r1.status_code == 200:
             return r1.text
         else:
             raise AttributeError("The daily list for this date is currently unavailable.")
 
+    def _download_file(self, list_id, full=False):
+        if full:
+            self._download_full_file(list_id)
+        else:
+            self._download_zip_file(list_id)
+        self._add_to_cache(list_id, full)
+
     def _download_zip_file(self, list_id):
         download_url = 'https://tranco-list.eu/download_daily/{}'.format(list_id)
         r = self.session.get(download_url, stream=True)
         if r.status_code == 200:
             with zipfile.ZipFile(BytesIO(r.content)) as z:
                 with z.open('top-1m.csv') as csvf:
                     file_bytes = csvf.read()
-                    if self.should_cache:
-                        with open(self._cache_path(list_id), 'wb') as f:
-                            f.write(file_bytes)
-                    lst = file_bytes.decode("utf-8")
-                    return lst
+                    with open(self._cache_path(list_id), 'wb') as f:
+                        f.write(file_bytes)
         elif r.status_code == 403:
             # List not available as ZIP file
             download_url = 'https://tranco-list.eu/download/{}/1000000'.format(list_id)
             r2 = self.session.get(download_url)
             if r2.status_code == 200:
                 file_bytes = r2.content
-                if self.should_cache:
-                    with open(self._cache_path(list_id), 'wb') as f:
-                        f.write(file_bytes)
-                lst = file_bytes.decode("utf-8")
-                return lst
+                with open(self._cache_path(list_id), 'wb') as f:
+                    f.write(file_bytes)
             else:
                 raise AttributeError("The daily list for this date is currently unavailable.")
         elif r.status_code == 502:
             # List unavailable (bad gateway)
-            raise AttributeError("The daily list for this date is currently unavailable.")
+            raise AttributeError("This list is currently unavailable.")
         else:
             # List unavailable (non-success status code)
             raise AttributeError("The daily list for this date is currently unavailable.")
 
+    def _download_full_file(self, list_id):
+        download_url = 'https://tranco-list.eu/download/{}/full'.format(list_id)
+        r = self.session.get(download_url)
+        if r.status_code == 200:
+            file_bytes = r.content
+            with open(self._cache_path(list_id), 'wb') as f:
+                f.write(file_bytes)
+
     def configure(self, configuration):
         """
         Configure a custom list (https://tranco-list.eu/configure).