feat: save downloaded packages to cache (#270)

Author: sdn4z

.gitignore

@@ -4,7 +4,7 @@
 *.pyc
 __pycache__
 db.sqlite3
-
+.twyn/
 # Backup files #
 *.bak
 

pyproject.toml

@@ -49,16 +49,18 @@ dev = [
     "pytest<9.0.0,>=7.1.3",
     "mypy<1.17,>=0.982",
     "pytest-cov<7,>=4",
-
     "ruff<0.12.4,>=0.5.1",
     "types-requests<3.0.0.0,>=2.32.4.20250611",
+    "types-python-dateutil>=2.9.0.20250809",
+     "freezegun>=1.5.5",
 ]
 local = [
     "ipdb<1.0.0,>=0.13.9",
     "commitizen<5.0,>=2.38",
     "pdbpp<1.0.0,>=0.11.6",
 ]
 
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

src/twyn/base/utils.py

@@ -0,0 +1,6 @@
+import re
+
+
+def _normalize_packages(packages: set[str]) -> set[str]:
+    """Normalize dependency names according to PyPi https://packaging.python.org/en/latest/specifications/name-normalization/."""
+    return {re.sub(r"[-_.]+", "-", name).lower() for name in packages}

src/twyn/cli.py

@@ -58,13 +58,20 @@ def entry_point() -> None:
     default=False,
     is_flag=True,
 )
+@click.option(
+    "--no-cache",
+    is_flag=True,
+    default=False,
+    help="Disable use of the trusted packages cache. Always fetch from the source.",
+)
 def run(
     config: str,
     dependency_file: Optional[str],
     dependency: tuple[str],
     selector_method: str,
     v: bool,
     vv: bool,
+    no_cache: bool,
 ) -> int:
     if v and vv:
         raise click.UsageError(
@@ -92,6 +99,7 @@ def run(
         dependency_file=dependency_file,
         selector_method=selector_method,
         verbosity=verbosity,
+        use_cache=not no_cache,
     )
 
     if errors:

src/twyn/dependency_parser/abstract_parser.py

@@ -26,7 +26,7 @@ def _read(self) -> str:
         return self.file_handler.read()
 
     def file_exists(self) -> bool:
-        return self.file_handler.file_exists()
+        return self.file_handler.exists()
 
     @abstractmethod
     def parse(self) -> set[str]:

src/twyn/file_handler/file_handler.py

@@ -11,7 +11,7 @@
 
 class BaseFileHandler(Protocol):
     def read(self) -> str: ...
-    def file_exists(self) -> bool: ...
+    def exists(self) -> bool: ...
     def write(self, data: str) -> None: ...
 
 
@@ -33,7 +33,7 @@ def read(self) -> str:
 
         return content
 
-    def file_exists(self) -> bool:
+    def exists(self) -> bool:
         try:
             self._raise_for_file_exists()
         except TwynError:
@@ -48,5 +48,4 @@ def _raise_for_file_exists(self) -> None:
             raise PathIsNotFileError
 
     def write(self, data: str) -> None:
-        self._raise_for_file_exists()
         self.file_path.write_text(data)

src/twyn/main.py

@@ -1,5 +1,4 @@
 import logging
-import re
 from typing import Optional
 
 from rich.logging import RichHandler
@@ -11,6 +10,7 @@
     AvailableLoggingLevels,
     SelectorMethod,
 )
+from twyn.base.utils import _normalize_packages
 from twyn.config.config_handler import ConfigHandler
 from twyn.dependency_parser.dependency_selector import DependencySelector
 from twyn.file_handler.file_handler import FileHandler
@@ -36,6 +36,7 @@ def check_dependencies(
     dependency_file: Optional[str] = None,
     dependencies: Optional[set[str]] = None,
     verbosity: AvailableLoggingLevels = AvailableLoggingLevels.none,
+    use_cache: bool = True,
 ) -> list[TyposquatCheckResult]:
     """Check if dependencies could be typosquats."""
     config_file_handler = FileHandler(config_file or DEFAULT_PROJECT_TOML_FILE)
@@ -45,7 +46,7 @@ def check_dependencies(
     _set_logging_level(config.logging_level)
 
     trusted_packages = TrustedPackages(
-        names=TopPyPiReference(source=config.pypi_reference).get_packages(),
+        names=TopPyPiReference(source=config.pypi_reference).get_packages(use_cache),
         algorithm=EditDistance(),
         selector=get_candidate_selector(config.selector_method),
         threshold_class=SimilarityThreshold,
@@ -84,8 +85,3 @@ def get_parsed_dependencies_from_file(dependency_file: Optional[str] = None) ->
     dependencies = dependency_parser.parse()
     logger.debug("Successfully parsed local dependencies file.")
     return dependencies
-
-
-def _normalize_packages(packages: set[str]) -> set[str]:
-    """Normalize dependency names according to PyPi https://packaging.python.org/en/latest/specifications/name-normalization/."""
-    return {re.sub(r"[-_.]+", "-", name).lower() for name in packages}

src/twyn/trusted_packages/constants.py

@@ -1,3 +1,8 @@
+# Cache configuration constants
+TRUSTED_PACKAGES_FILE_PATH = ".twyn/trusted_packages.json"
+TRUSTED_PACKAGES_MAX_RETENTION_DAYS = 30
+
+
 ADJACENCY_MATRIX = {
     "1": ["2", "q", "w"],
     "2": ["1", "3", "q", "w"],

src/twyn/trusted_packages/exceptions.py

@@ -13,4 +13,8 @@ class EmptyPackagesListError(TwynError):
     message = "Downloaded packages list is empty"
 
 
-class CharacterNotInMatrixError(TwynError, KeyError): ...
+class CharacterNotInMatrixError(TwynError, KeyError): ...  # TODO add comments
+
+
+class InvalidCacheError(TwynError):
+    """Error for when the cache content is not valid."""

src/twyn/trusted_packages/references.py

@@ -1,11 +1,17 @@
+import json
 import logging
 from abc import ABC, abstractmethod
+from datetime import date, datetime
 from typing import Any
 
 import requests
 
+from twyn.base.utils import _normalize_packages
+from twyn.file_handler.file_handler import FileHandler
+from twyn.trusted_packages.constants import TRUSTED_PACKAGES_FILE_PATH, TRUSTED_PACKAGES_MAX_RETENTION_DAYS
 from twyn.trusted_packages.exceptions import (
     EmptyPackagesListError,
+    InvalidCacheError,
     InvalidJSONError,
     InvalidPyPiFormatError,
 )
@@ -20,17 +26,95 @@ def __init__(self, source: str) -> None:
         self.source = source
 
     @abstractmethod
-    def get_packages(self) -> set[str]:
+    def get_packages(self, use_cache: bool = True) -> set[str]:
         """Return the names of the trusted packages available in the reference."""
 
 
 class TopPyPiReference(AbstractPackageReference):
     """Top PyPi packages retrieved from an online source."""
 
-    def get_packages(self) -> set[str]:
+    def get_packages(self, use_cache: bool = True) -> set[str]:
         """Download and parse online source of top Python Package Index packages."""
-        packages_info = self._download()
-        return self._parse(packages_info)
+        packages_to_use = set()
+        if use_cache:
+            trusted_packages_file = FileHandler(TRUSTED_PACKAGES_FILE_PATH)
+            packages_to_use = self._get_packages_from_cache(trusted_packages_file)
+            # we don't save the cache here, we keep it as it is so the date remains the original one.
+
+        if not packages_to_use:
+            # no cache usage, no cache hit (non-existent or outdated) or cache was empty.
+            logger.info("Fetching trusted packages from PyPI reference...")
+            packages_to_use = self._parse(self._download())
+            if use_cache:
+                self._save_trusted_packages_to_file(packages_to_use, trusted_packages_file, self.source)
+
+        normalized_packages = _normalize_packages(packages_to_use)
+        return normalized_packages
+
+    def _is_content_outdated(self, content_date: date) -> bool:
+        """Check if cached content is outdated based on retention days."""
+        days_diff = (datetime.today().date() - content_date).days
+        return days_diff > TRUSTED_PACKAGES_MAX_RETENTION_DAYS
+
+    def _save_trusted_packages_to_file(self, packages: set[str], file_handler: FileHandler, source: str) -> None:
+        """Save trusted packages to JSON file with timestamp."""
+        trusted_data = {
+            "source": source,
+            "data": {
+                "packages": list(packages),
+                "count": len(packages),
+                "saved_date": datetime.now().date().isoformat(),
+            },
+        }
+        file_handler.file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_handler.write(json.dumps(trusted_data))
+        logger.debug("Saved %d trusted packages to %s", len(packages), file_handler.file_path)
+
+    def _load_trusted_packages_from_file(self, file_handler: FileHandler) -> tuple[set[str], bool]:
+        """Load trusted packages from JSON file and check if it's outdated."""
+        try:
+            try:
+                trusted_packages_raw_content = json.loads(file_handler.read())
+            except json.JSONDecodeError as e:
+                raise InvalidCacheError("Could not decode cache.") from e
+
+            try:
+                data = trusted_packages_raw_content["data"]
+                saved_date_str = data["saved_date"]
+            except KeyError as e:
+                raise InvalidCacheError("Invalid cache format.") from e
+
+            try:
+                saved_date = datetime.fromisoformat(saved_date_str).date()
+            except ValueError as e:
+                raise InvalidCacheError("Cache saved date is invalid.") from e
+
+            try:
+                packages = set(data["packages"])
+            except TypeError as e:
+                raise InvalidCacheError("Invalid format in cached packages") from e
+
+            is_outdated = self._is_content_outdated(saved_date)
+
+        except InvalidCacheError as e:
+            logger.warning("Error reading cached trusted packages: %s", e)
+            return set(), True
+        else:
+            if is_outdated:
+                logger.info("Cached trusted packages are outdated (saved: %s)", saved_date)
+            else:
+                logger.debug("Using cached trusted packages from %s", saved_date)
+
+            return packages, is_outdated
+
+    def _get_packages_from_cache(self, trusted_packages_file: FileHandler) -> set[str]:
+        """Get packages from cache file if it's present and up to date."""
+        if trusted_packages_file.exists():
+            packages_from_cache, is_outdated = self._load_trusted_packages_from_file(trusted_packages_file)
+            if not is_outdated:
+                return packages_from_cache
+
+        return set()
 
     def _download(self) -> dict[str, Any]:
         packages = requests.get(self.source)