Requirement version resolution fix

mlem/core/hooks.py

@@ -67,7 +67,7 @@ def __init_subclass__(cls, *args, **kwargs):
                 )
         else:
             logger.debug(
-                "Not registerting %s to any Analyzer because it's an abstract class",
+                "Not registering %s to any Analyzer because it's an abstract class",
                 cls.__name__,
             )
         super(Hook, cls).__init_subclass__(*args, **kwargs)

mlem/utils/module.py

@@ -10,6 +10,7 @@
 import warnings
 from collections import defaultdict
 from functools import lru_cache, wraps
+from importlib.metadata import PackageNotFoundError, distribution
 from pickle import PickleError
 from types import FunctionType, LambdaType, MethodType, ModuleType
 from typing import Dict, List, Optional, Set, Union
@@ -249,15 +250,22 @@ def get_module_version(mod: ModuleType) -> Optional[str]:
     :param mod: module object to use
     :return: version as `str` or `None` if version could not be determined
     """
-    for attr in "__version__", "VERSION":
-        if hasattr(mod, attr):
-            return str(getattr(mod, attr))
-    if mod.__file__ is None:
-        return None
-    for name in os.listdir(os.path.dirname(mod.__file__)):
-        m = re.match(re.escape(mod.__name__) + "-(.+)\\.dist-info", name)
-        if m:
-            return m.group(1)
+    # try using importlib package -> distro mapping and distro metadata
+    package_to_distros = packages_distributions()
+    try:
+        if mod.__name__ in package_to_distros:
+            for distro_name in package_to_distros[mod.__name__]:
+                distro = distribution(distro_name)
+                return distro.version
+    except PackageNotFoundError:
+        pass
+
+    # if there's a package-file, try to get it from there
+    if mod.__file__ is not None:  # pragma: no branch
+        for name in os.listdir(os.path.dirname(mod.__file__)):
+            m = re.match(re.escape(mod.__name__) + "-(.+)\\.dist-info", name)
+            if m:
+                return m.group(1)
     return None
 
 
@@ -502,7 +510,7 @@ def save_type_with_classvars(pickler: "RequirementAnalyzer", obj):
 
 class RequirementAnalyzer(dill.Pickler):
     """Special pickler implementation that collects requirements while pickling
-    (and not pickling actualy)"""
+    (and not actually pickling)"""
 
     ignoring = (
         "dill",

setup.py

@@ -57,6 +57,9 @@
     "jupyter",
     "nbconvert",
     "nbloader",
+    # We're using regex to test requirement version extraction
+    # edge case, see: https://github.com/iterative/mlem/issues/688
+    "regex==2023.6.3",
 ]
 
 extras = {

tests/utils/test_module_tools.py

@@ -5,6 +5,7 @@
 import nbformat
 import numpy
 import pytest
+import regex
 from pydantic import BaseModel
 
 from mlem.core.requirements import Requirements
@@ -220,11 +221,11 @@ def test_get_requirements_notebook():
 
     loaded_notebook = Notebook(TEST_NOTEBOOK_NAME)
 
-    kek = loaded_notebook.run_all()
-    res = kek.ns["res"]
+    notebook = loaded_notebook.run_all()
+    res = notebook.ns["res"]
 
     assert isinstance(res, Requirements)
-    assert res.modules == ["numpy"]
+    assert sorted(res.modules) == sorted(["regex", "numpy"])
 
 
 def _run_jup(command):
@@ -249,7 +250,9 @@ def test_get_requirements_notebook_run():
     with open(TEST_NOTEBOOK_NAME, encoding="utf8") as f:
         nb = nbformat.read(f, as_version=4)
 
-    assert nb["cells"][1]["outputs"][0].text.strip() == get_module_repr(numpy)
+    res = nb["cells"][1]["outputs"][0].text.strip()
+    assert get_module_repr(numpy) in res
+    assert get_module_repr(regex) in res
 
 
 # Copyright 2019 Zyfra

tests/utils/test_save.ipynb

@@ -5,10 +5,10 @@
    "execution_count": 1,
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2023-02-13T14:11:29.261665Z",
-     "iopub.status.busy": "2023-02-13T14:11:29.261394Z",
-     "iopub.status.idle": "2023-02-13T14:11:29.267566Z",
-     "shell.execute_reply": "2023-02-13T14:11:29.266734Z"
+     "iopub.execute_input": "2023-07-31T10:02:34.099802Z",
+     "iopub.status.busy": "2023-07-31T10:02:34.099406Z",
+     "iopub.status.idle": "2023-07-31T10:02:34.115259Z",
+     "shell.execute_reply": "2023-07-31T10:02:34.114688Z"
     },
     "pycharm": {
      "name": "#%%\n"
@@ -17,8 +17,11 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
+    "import regex\n",
     "\n",
     "def func(data):\n",
+    "    # using regex, just so it's listed and its version inferred correctly\n",
+    "    regex.search(r'(ab)', 'abcdef')\n",
     "    return bool(np.all([True]))\n"
    ]
   },
@@ -27,10 +30,10 @@
    "execution_count": 2,
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2023-02-13T14:11:29.270506Z",
-     "iopub.status.busy": "2023-02-13T14:11:29.270233Z",
-     "iopub.status.idle": "2023-02-13T14:11:29.969468Z",
-     "shell.execute_reply": "2023-02-13T14:11:29.968705Z"
+     "iopub.execute_input": "2023-07-31T10:02:34.118394Z",
+     "iopub.status.busy": "2023-07-31T10:02:34.118192Z",
+     "iopub.status.idle": "2023-07-31T10:02:35.897692Z",
+     "shell.execute_reply": "2023-07-31T10:02:35.897269Z"
     },
     "pycharm": {
      "is_executing": true,
@@ -42,7 +45,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "numpy==1.22.4\n"
+      "numpy==1.25.1 regex==2023.6.3\n"
      ]
     }
    ],
@@ -71,7 +74,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.13"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,