From 2c2c5098f0e97027c84989e49e738851293b605b Mon Sep 17 00:00:00 2001
From: Saaket Prakash <saaketp@gmail.com>
Date: Thu, 8 Sep 2022 12:02:38 +0530
Subject: [PATCH] Use importlib.metadata for guess_version instead of importing
 the module (#171)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Bernát Gábor <bgabor8@bloomberg.net>
---
 CHANGES.md                        |  4 ++++
 pyproject.toml                    |  2 +-
 src/pipdeptree/__init__.py        | 11 +++++++++++
 tests/guess_version_setuptools.py | 16 ++++++++++++++++
 tests/test_pipdeptree.py          |  8 ++++++++
 5 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 tests/guess_version_setuptools.py

diff --git a/CHANGES.md b/CHANGES.md
index cd945dd..23fd608 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## 2.3.1
+
+- Use `importlib.metadata` to guess version of package before fallback to `pkg.__version__`.
+
 ## 2.3.0
 
 - Move to a package layout
diff --git a/pyproject.toml b/pyproject.toml
index d3862a2..fd5f2d6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,7 @@ optional-dependencies.test = [
   "pytest>=7.1.3",
   "pytest-cov>=3",
   "pytest-mock>=3.8.2",
-  "virtualenv<21,>=20.16.4",
+  "virtualenv<21,>=20.16.5",
 ]
 optional-dependencies.graphviz = [
   "graphviz>=0.20.1",
diff --git a/src/pipdeptree/__init__.py b/src/pipdeptree/__init__.py
index 9040088..45d13b4 100644
--- a/src/pipdeptree/__init__.py
+++ b/src/pipdeptree/__init__.py
@@ -41,6 +41,17 @@ def guess_version(pkg_key, default="?"):
     :returns: version
     :rtype: string
     """
+    try:
+        if sys.version_info >= (3, 8):  # pragma: >=3.8 cover
+            import importlib.metadata as importlib_metadata
+        else:  # pragma: <3.8 cover
+            import importlib_metadata
+        return importlib_metadata.version(pkg_key)
+    except ImportError:
+        pass
+    # Avoid AssertionError with setuptools, see https://github.com/tox-dev/pipdeptree/issues/162
+    if pkg_key in {"setuptools"}:
+        return default
     try:
         m = import_module(pkg_key)
     except ImportError:
diff --git a/tests/guess_version_setuptools.py b/tests/guess_version_setuptools.py
new file mode 100644
index 0000000..f9426f7
--- /dev/null
+++ b/tests/guess_version_setuptools.py
@@ -0,0 +1,16 @@
+import sys
+
+import pipdeptree
+
+if sys.version_info >= (3, 8):
+    import importlib.metadata as importlib_metadata
+else:
+    import importlib_metadata
+
+
+def raise_import_error(name):
+    raise ImportError(name)
+
+
+importlib_metadata.version = raise_import_error
+print(pipdeptree.guess_version("setuptools"), end="")
diff --git a/tests/test_pipdeptree.py b/tests/test_pipdeptree.py
index 6bde2fd..27c77dc 100644
--- a/tests/test_pipdeptree.py
+++ b/tests/test_pipdeptree.py
@@ -1,7 +1,9 @@
 import platform
+import subprocess
 import sys
 from contextlib import contextmanager
 from itertools import chain
+from pathlib import Path
 from tempfile import NamedTemporaryFile
 
 try:
@@ -488,3 +490,9 @@ def test_custom_interpreter(tmp_path, monkeypatch, capfd, args_joined):
     assert context.value.code == 1
     assert not out
     assert err == "graphviz functionality is not supported when querying" " non-host python\n"
+
+
+def test_guess_version_setuptools():
+    script = Path(__file__).parent / "guess_version_setuptools.py"
+    output = subprocess.check_output([sys.executable, script], text=True)
+    assert output == "?"