Add requires_package decorator

Author: vhaldemar

src/yandex_cloud_ml_sdk/_utils/packages.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import functools
+import importlib.metadata
+import inspect
+import types
+from typing import Callable, TypeVar
+
+import packaging.specifiers
+import packaging.version
+from typing_extensions import ParamSpec
+
+P = ParamSpec('P')
+R = TypeVar('R')
+
+
+def requires_package(
+    package_name: str, requirements: str, method_name: str
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    required_versions = packaging.specifiers.SpecifierSet(requirements)
+
+    def decorator(method: Callable[P, R]) -> Callable[P, R]:
+        @functools.wraps(method)
+        def function(*args: P.args, **kwargs: P.kwargs) -> R:
+            error: str | None = None
+            try:
+                package_version = importlib.metadata.version(package_name)
+                installed_version = packaging.version.parse(package_version)
+
+                if installed_version not in required_versions:
+                    error = f', but you have "{installed_version}" installed'
+            except importlib.metadata.PackageNotFoundError:
+                error = ", but you don't have any version installed"
+            if error is not None:
+                e = RuntimeError(
+                    f'"{method_name}" requires package "{package_name}{requirements}" to be installed{error}'
+                )
+                current_frame = inspect.currentframe()
+                new_tb: types.TracebackType | None = None
+                if current_frame and current_frame.f_back:
+                    caller_frame = current_frame.f_back
+                    new_tb = types.TracebackType(
+                        tb_next=None,
+                        tb_frame=caller_frame,
+                        tb_lasti=caller_frame.f_lasti,
+                        tb_lineno=caller_frame.f_lineno,
+                    )
+                raise e.with_traceback(new_tb)
+
+            return method(*args, **kwargs)
+        return function
+    return decorator

tests/utils/test_packages.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import importlib.metadata
+
+import pytest
+
+from yandex_cloud_ml_sdk._utils.packages import requires_package
+
+
+@requires_package('somerandomname', '>1', 'mysupermethod')
+def mysupermethod():
+    pass
+
+
+@requires_package('yandexcloud', '<0.100', 'mysupermethod')
+def mysupermethod2():
+    pass
+
+
+@requires_package('yandexcloud', '>0.100', 'mysupermethod')
+def mysupermethod3():
+    pass
+
+
+def test_requires_package():
+    with pytest.raises(RuntimeError) as exc_info:
+        mysupermethod()
+    assert exc_info.value.args[0] == (
+        '"mysupermethod" requires package "somerandomname>1"'
+        ' to be installed, but you don\'t have any version installed'
+    )
+    assert str(exc_info.traceback[0].path) == __file__
+    assert str(exc_info.traceback[-1].path) == __file__
+
+    installed_version = importlib.metadata.version('yandexcloud')
+    with pytest.raises(
+        RuntimeError,
+        match=(
+            '"mysupermethod" requires package "yandexcloud<0.100" to be installed, '
+            f'but you have "{installed_version}" installed'
+        )
+    ):
+        mysupermethod2()
+
+    mysupermethod3()