Skip to content

Commit 2af3238

Browse files
fix: replace pkg_resources in system parameter logging
Collect installed package versions with importlib.metadata and use the importlib-metadata backport on Python 3.7. Preserve legacy package-name normalization, cache the distribution scan, and return copies so callers cannot modify cached data. Add regression coverage.
1 parent 6e098e3 commit 2af3238

3 files changed

Lines changed: 70 additions & 4 deletions

File tree

aim/ext/utils.py

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,52 @@
11
import logging
2+
import re
23
import subprocess
34

5+
from functools import lru_cache
6+
47
from fastapi.responses import JSONResponse
58

69

710
logger = logging.getLogger(__name__)
811

912

10-
def get_installed_packages():
11-
import pkg_resources
13+
def _get_installed_distributions():
14+
try:
15+
from importlib import metadata as metadata_module
16+
except ImportError:
17+
import importlib_metadata as metadata_module # Python 3.7 support
18+
19+
return metadata_module.distributions()
1220

13-
packages = {i.key: i.version for i in pkg_resources.working_set}
1421

15-
return packages
22+
# Installed distributions normally do not change while the current Python process is running.
23+
@lru_cache(maxsize=1)
24+
def _collect_installed_package_versions():
25+
installed_package_versions = {}
26+
27+
for installed_distribution in _get_installed_distributions():
28+
distribution_metadata = installed_distribution.metadata
29+
package_name = distribution_metadata.get('Name')
30+
package_version = distribution_metadata.get('Version')
31+
32+
if not package_name or not package_version:
33+
continue
34+
35+
# Preserve the normalized key format previously provided by pkg_resources.
36+
normalized_package_name = re.sub(
37+
r'[^A-Za-z0-9.]+',
38+
'-',
39+
package_name,
40+
).lower()
41+
42+
installed_package_versions.setdefault(normalized_package_name, package_version)
43+
44+
return installed_package_versions
45+
46+
47+
def get_installed_packages():
48+
# Return a copy so callers cannot modify the cached package information.
49+
return _collect_installed_package_versions().copy()
1650

1751

1852
def get_environment_variables():

setup.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ def package_files(directory):
7474
'uvicorn<1,>=0.12.0',
7575
'Pillow>=8.0.0',
7676
'packaging>=15.0',
77+
'importlib-metadata<6.8; python_version < "3.8"',
7778
'python-dateutil',
7879
'requests',
7980
'watchdog',

tests/ext/test_utils.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import unittest
2+
3+
from types import SimpleNamespace
4+
from unittest import mock
5+
6+
from aim.ext import utils
7+
8+
9+
class TestInstalledPackages(unittest.TestCase):
10+
def setUp(self):
11+
utils._collect_installed_package_versions.cache_clear()
12+
13+
def tearDown(self):
14+
utils._collect_installed_package_versions.cache_clear()
15+
16+
@mock.patch('aim.ext.utils._get_installed_distributions')
17+
def test_get_installed_packages(self, mock_installed_distributions):
18+
mock_installed_distributions.return_value = [
19+
SimpleNamespace(metadata={'Name': 'Example_Package', 'Version': '1.2.3'}),
20+
SimpleNamespace(metadata={'Version': '2.0.0'}),
21+
SimpleNamespace(metadata={'Name': 'missing-version'}),
22+
]
23+
24+
packages = utils.get_installed_packages()
25+
packages['changed-by-caller'] = '1.0'
26+
27+
self.assertEqual(
28+
{'example-package': '1.2.3'},
29+
utils.get_installed_packages(),
30+
)
31+
mock_installed_distributions.assert_called_once_with()

0 commit comments

Comments
 (0)