From a2c03eb4fb78f8b2f6ca4a07e0dbae79bc69fac2 Mon Sep 17 00:00:00 2001
From: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Date: Fri, 22 Mar 2024 14:26:59 +0100
Subject: [PATCH] docs: fix/use PyPI versions for pinning (#243)

* docs: fix/use PyPI versions for pinning
* chlog & fixing
---
 CHANGELOG.md                               |  2 +-
 src/lightning_utilities/docs/formatting.py | 34 ++++++++++++++++++----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index db3cda1a..620d3a8b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,7 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
--
+- docs: fix/use PyPI versions for pinning links ([#243](https://github.com/Lightning-AI/utilities/pull/243))
 
 
 ## [0.11.0] - 2024-03-18
diff --git a/src/lightning_utilities/docs/formatting.py b/src/lightning_utilities/docs/formatting.py
index 32512ca4..32f616ac 100644
--- a/src/lightning_utilities/docs/formatting.py
+++ b/src/lightning_utilities/docs/formatting.py
@@ -8,7 +8,7 @@
 import os
 import re
 import sys
-from typing import Iterable, Optional, Tuple, Union
+from typing import Iterable, List, Optional, Tuple, Union
 
 
 def _transform_changelog(path_in: str, path_out: str) -> None:
@@ -69,6 +69,27 @@ def find_source() -> Tuple[str, int, int]:
     return f"https://github.com/{github_user}/{github_repo}/blob/{filename}"
 
 
+def _load_pypi_versions(package_name: str) -> List[str]:
+    """Load the versions of the package from PyPI.
+
+    >>> _load_pypi_versions("numpy")  # doctest: +ELLIPSIS
+    ['0.9.6', '0.9.8', '1.0', ...]
+    >>> _load_pypi_versions("scikit-learn")  # doctest: +ELLIPSIS
+    ['0.9', '0.10', '0.11', '0.12', ...]
+
+    """
+    from distutils.version import LooseVersion
+
+    import requests
+
+    url = f"https://pypi.org/pypi/{package_name}/json"
+    data = requests.get(url, timeout=10).json()
+    versions = data["releases"].keys()
+    # filter all version which include only numbers and dots
+    versions = {k for k in versions if re.match(r"^\d+(\.\d+)*$", k)}
+    return sorted(versions, key=LooseVersion)
+
+
 def _update_link_based_imported_package(link: str, pkg_ver: str, version_digits: Optional[int]) -> str:
     """Adjust the linked external docs to be local.
 
@@ -79,10 +100,13 @@ def _update_link_based_imported_package(link: str, pkg_ver: str, version_digits:
 
     """
     pkg_att = pkg_ver.split(".")
-    # load the package with all additional sub-modules
-    module = importlib.import_module(".".join(pkg_att[:-1]))
-    # load the attribute
-    ver = getattr(module, pkg_att[-1])
+    try:
+        ver = _load_pypi_versions(pkg_att[0])[-1]
+    except Exception:
+        # load the package with all additional sub-modules
+        module = importlib.import_module(".".join(pkg_att[:-1]))
+        # load the attribute
+        ver = getattr(module, pkg_att[0])
     # drop any additional context after `+`
     ver = ver.split("+")[0]
     # crop the version to the number of digits