From 678b512aed785243ffa169aa3feab2e57cd4cd40 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Thu, 13 Mar 2025 14:24:55 -0400
Subject: [PATCH] ci: limit downloads from pypi stats to main branch (#3835)

At some point, we can run on a cron schedule
---
 .github/workflows/deploy_docs.yml             |  9 +++++-
 docs/Makefile                                 | 10 ++++++-
 .../third_party_page/get_download_stats.py    | 30 ++++++++++++++++---
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/deploy_docs.yml b/.github/workflows/deploy_docs.yml
index e7ea2cc1e..0c9c77590 100644
--- a/.github/workflows/deploy_docs.yml
+++ b/.github/workflows/deploy_docs.yml
@@ -102,7 +102,14 @@ jobs:
       - name: Build llms-text
         run: make llms-text
       - name: Build site
-        run: make build-docs
+        run: |
+          # If this is main branch, then we want to download stats. we do this 
+          # with the env variable DOWNLOAD_STATS=true
+          if [ "${{ github.ref }}" == "refs/heads/main" ]; then
+            DOWNLOAD_STATS=true make build-docs
+          else
+            make build-docs
+          fi
         env:
           MKDOCS_GIT_COMMITTERS_APIKEY: ${{ secrets.MKDOCS_GIT_COMMITTERS_APIKEY }}
           OPENAI_API_KEY: sf-proj-1234567890 # fake placeholder, shouldn't actually be used
diff --git a/docs/Makefile b/docs/Makefile
index d8df63efd..487502fc4 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -10,7 +10,15 @@ build-prebuilt:
 	# Use to create an update to date prebuilt page.
 	# Looks up download stats for each of the prebuilt packages and
 	# generates the final prebuilt page.
-	poetry run python -m _scripts.third_party_page.get_download_stats stats.yml
+	@if [ "$(DOWNLOAD_STATS)" = "true" ]; then \
+		set -x; \
+		poetry run python -m _scripts.third_party_page.get_download_stats stats.yml; \
+		set +x; \
+	else \
+		set -x; \
+		poetry run python -m _scripts.third_party_page.get_download_stats --fake stats.yml; \
+		set +x; \
+	fi
 	poetry run python -m _scripts.third_party_page.create_third_party_page stats.yml docs/prebuilt.md --language python
 
 build-docs: build-typedoc build-prebuilt
diff --git a/docs/_scripts/third_party_page/get_download_stats.py b/docs/_scripts/third_party_page/get_download_stats.py
index 81b6b515f..582aa269f 100755
--- a/docs/_scripts/third_party_page/get_download_stats.py
+++ b/docs/_scripts/third_party_page/get_download_stats.py
@@ -30,10 +30,23 @@ PACKAGES_FILE = HERE / "packages.yml"
 PACKAGES = yaml.safe_load(PACKAGES_FILE.read_text())['packages']
 
 
-def _get_weekly_downloads(packages: list[Package]) -> list[ResolvedPackage]:
+def _get_weekly_downloads(packages: list[Package], fake: bool) -> list[ResolvedPackage]:
     """Retrieve the monthly download count for a list of packages from PyPIStats."""
     resolved_packages: list[ResolvedPackage] = []
 
+    if fake:
+        # To avoid making network requests during testing, return fake download counts
+        for package in packages:
+            resolved_packages.append(
+                {
+                    "name": package["name"],
+                    "repo": package["repo"],
+                    "weekly_downloads": -12345,
+                    "description": package["description"],
+                }
+            )
+        return resolved_packages
+
     for package in packages:
         # First check if package exists on PyPI
         pypi_url = f"https://pypi.org/pypi/{package['name']}/json"
@@ -88,13 +101,13 @@ def _get_weekly_downloads(packages: list[Package]) -> list[ResolvedPackage]:
 
 
 
-def main(output_file: str) -> None:
+def main(output_file: str, fake: bool) -> None:
     """Main function to generate package download information.
 
     Args:
         output_file: Path to the output YAML file.
     """
-    resolved_packages: list[ResolvedPackage] = _get_weekly_downloads(PACKAGES)
+    resolved_packages: list[ResolvedPackage] = _get_weekly_downloads(PACKAGES, fake)
 
     if not output_file.endswith(".yml"):
         raise ValueError("Output file must have a .yml extension")
@@ -115,6 +128,15 @@ if __name__ == "__main__":
             "downloads.yml"
         ),
     )
+    parser.add_argument(
+        "--fake",
+        default=False,
+        action="store_true",
+        help=(
+            "Generate fake download counts for testing purposes. "
+            "This option will not make any network requests."
+        ),
+    )
     args = parser.parse_args()
 
-    main(args.output_file)
+    main(args.output_file, args.fake)