From bb820ab3b085ca28af9a6f839fac60d491ac2168 Mon Sep 17 00:00:00 2001
From: awaelchli <aedu.waelchli@gmail.com>
Date: Sat, 13 Jul 2024 09:45:48 +0200
Subject: [PATCH] Replace deprecated `pkg_resources` (#281)

* replace deprecated pkg_resources
* chlog

---------

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
---
 CHANGELOG.md                            |  9 +++++++
 src/lightning_utilities/core/imports.py | 31 +++++++++++++------------
 tests/unittests/core/test_imports.py    |  2 +-
 3 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e0dbb985..de1702ba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ---
+## [UnRaleased] - 2024-MM-DD
+
+### Changed
+
+- Replaced deprecated `pkg_resources` with `importlib.metadata` ([#281](https://github.com/Lightning-AI/utilities/pull/281))
+
+### Fixed
+
+-
 
 ## [0.11.3] - 2024-06-26
 
diff --git a/src/lightning_utilities/core/imports.py b/src/lightning_utilities/core/imports.py
index 461ea934..5362845b 100644
--- a/src/lightning_utilities/core/imports.py
+++ b/src/lightning_utilities/core/imports.py
@@ -7,13 +7,14 @@
 import os
 import warnings
 from functools import lru_cache
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
 from importlib.util import find_spec
 from types import ModuleType
 from typing import Any, Callable, List, Optional, TypeVar
 
-import pkg_resources
 from packaging.requirements import Requirement
-from packaging.version import Version
+from packaging.version import InvalidVersion, Version
 from typing_extensions import ParamSpec
 
 T = TypeVar("T")
@@ -75,15 +76,12 @@ def compare_version(package: str, op: Callable, version: str, use_base_version:
     """
     try:
         pkg = importlib.import_module(package)
-    except (ImportError, pkg_resources.DistributionNotFound):
+    except ImportError:
         return False
     try:
-        if hasattr(pkg, "__version__"):
-            pkg_version = Version(pkg.__version__)
-        else:
-            # try pkg_resources to infer version
-            pkg_version = Version(pkg_resources.get_distribution(package).version)
-    except TypeError:
+        # Use importlib.metadata to infer version
+        pkg_version = Version(pkg.__version__) if hasattr(pkg, "__version__") else Version(_version(package))
+    except (TypeError, PackageNotFoundError):
         # this is mocked by Sphinx, so it should return True to generate all summaries
         return True
     if use_base_version:
@@ -128,17 +126,20 @@ def __init__(self, requirement: Optional[str] = None, module: Optional[str] = No
     def _check_requirement(self) -> None:
         assert self.requirement  # noqa: S101; needed for typing
         try:
-            # first try the pkg_resources requirement
-            pkg_resources.require(self.requirement)
-            self.available = True
-            self.message = f"Requirement {self.requirement!r} met"
-        except Exception as ex:
+            req = Requirement(self.requirement)
+            pkg_version = Version(_version(req.name))
+            self.available = req.specifier.contains(pkg_version)
+        except (PackageNotFoundError, InvalidVersion) as ex:
             self.available = False
             self.message = f"{ex.__class__.__name__}: {ex}. HINT: Try running `pip install -U {self.requirement!r}`"
+
+        if self.available:
+            self.message = f"Requirement {self.requirement!r} met"
+        else:
             req_include_version = any(c in self.requirement for c in "=<>")
             if not req_include_version or self.module is not None:
                 module = self.requirement if self.module is None else self.module
-                # sometimes `pkg_resources.require()` fails but the module is importable
+                # Sometimes `importlib.metadata.version` fails but the module is importable
                 self.available = module_available(module)
                 if self.available:
                     self.message = f"Module {module!r} available"
diff --git a/tests/unittests/core/test_imports.py b/tests/unittests/core/test_imports.py
index e25183a7..7e3998f9 100644
--- a/tests/unittests/core/test_imports.py
+++ b/tests/unittests/core/test_imports.py
@@ -48,7 +48,7 @@ def testcompare_version(monkeypatch):
 def test_requirement_cache():
     assert RequirementCache(f"pytest>={pytest.__version__}")
     assert not RequirementCache(f"pytest<{pytest.__version__}")
-    assert "pip install -U '-'" in str(RequirementCache("-"))
+    assert "pip install -U 'not-found-requirement'" in str(RequirementCache("not-found-requirement"))
 
     # invalid requirement is skipped by valid module
     assert RequirementCache(f"pytest<{pytest.__version__}", "pytest")