Switch from pkg_resources to importlib (#3722)

* Switch from pkg_resources to importlib

* Some test fixes

* Improve marker parsing

* WIP databricks

* Databricks fixes

* Revert some changes

* Linting

* Fix import

---------

Co-authored-by: Hamza Tahir <[email protected]>

Author: schustmi

src/zenml/cli/base.py

@@ -49,6 +49,7 @@
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils.io_utils import copy_dir, get_global_config_directory
+from zenml.utils.package_utils import get_package_information
 from zenml.utils.server_utils import get_local_server
 from zenml.utils.yaml_utils import write_yaml
 
@@ -640,7 +641,7 @@ def info(
     }
 
     if all:
-        user_info["packages"] = cli_utils.get_package_information()
+        user_info["packages"] = get_package_information()
     if packages:
         if user_info.get("packages"):
             if isinstance(user_info["packages"], dict):
@@ -650,7 +651,7 @@ def info(
                     if p in packages
                 }
         else:
-            user_info["query_packages"] = cli_utils.get_package_information(
+            user_info["query_packages"] = get_package_information(
                 list(packages)
             )
     if file:

src/zenml/cli/utils.py

@@ -40,7 +40,6 @@
 )
 
 import click
-import pkg_resources
 import yaml
 from pydantic import BaseModel, SecretStr
 from rich import box, table
@@ -80,6 +79,7 @@
 from zenml.stack.flavor import Flavor
 from zenml.stack.stack_component import StackComponentConfig
 from zenml.utils import secret_utils
+from zenml.utils.package_utils import requirement_installed
 from zenml.utils.time_utils import expires_in
 from zenml.utils.typing_utils import get_origin, is_union
 
@@ -1052,7 +1052,7 @@ def install_packages(
         # just return without doing anything
         return
 
-    if use_uv and not is_installed_in_python_environment("uv"):
+    if use_uv and not requirement_installed("uv"):
         # If uv is installed globally, don't run as a python module
         command = []
     else:
@@ -1094,7 +1094,7 @@ def uninstall_package(package: str, use_uv: bool = False) -> None:
         package: The package to uninstall.
         use_uv: Whether to use uv for package uninstallation.
     """
-    if use_uv and not is_installed_in_python_environment("uv"):
+    if use_uv and not requirement_installed("uv"):
         # If uv is installed globally, don't run as a python module
         command = []
     else:
@@ -1110,22 +1110,6 @@ def uninstall_package(package: str, use_uv: bool = False) -> None:
     subprocess.check_call(command)
 
 
-def is_installed_in_python_environment(package: str) -> bool:
-    """Check if a package is installed in the current python environment.
-
-    Args:
-        package: The package to check.
-
-    Returns:
-        True if the package is installed, False otherwise.
-    """
-    try:
-        pkg_resources.get_distribution(package)
-        return True
-    except pkg_resources.DistributionNotFound:
-        return False
-
-
 def is_uv_installed() -> bool:
     """Check if uv is installed.
 
@@ -1141,7 +1125,7 @@ def is_pip_installed() -> bool:
     Returns:
         True if pip is installed, False otherwise.
     """
-    return is_installed_in_python_environment("pip")
+    return requirement_installed("pip")
 
 
 def pretty_print_secret(
@@ -2499,30 +2483,6 @@ def temporary_active_stack(
             Client().activate_stack(old_stack_id)
 
 
-def get_package_information(
-    package_names: Optional[List[str]] = None,
-) -> Dict[str, str]:
-    """Get a dictionary of installed packages.
-
-    Args:
-        package_names: Specific package names to get the information for.
-
-    Returns:
-        A dictionary of the name:version for the package names passed in or
-            all packages and their respective versions.
-    """
-    import pkg_resources
-
-    if package_names:
-        return {
-            pkg.key: pkg.version
-            for pkg in pkg_resources.working_set
-            if pkg.key in package_names
-        }
-
-    return {pkg.key: pkg.version for pkg in pkg_resources.working_set}
-
-
 def print_user_info(info: Dict[str, Any]) -> None:
     """Print user information to the terminal.
 
@@ -2617,11 +2577,7 @@ def is_jupyter_installed() -> bool:
     Returns:
         bool: True if Jupyter notebook is installed, False otherwise.
     """
-    try:
-        pkg_resources.get_distribution("notebook")
-        return True
-    except pkg_resources.DistributionNotFound:
-        return False
+    return requirement_installed("notebook")
 
 
 def multi_choice_prompt(

src/zenml/integrations/databricks/orchestrators/databricks_orchestrator_entrypoint_config.py

@@ -17,7 +17,10 @@
 import sys
 from typing import Any, List, Set
 
-import pkg_resources
+if sys.version_info < (3, 10):
+    from importlib_metadata import distribution
+else:
+    from importlib.metadata import distribution
 
 from zenml.entrypoints.step_entrypoint_configuration import (
     StepEntrypointConfiguration,
@@ -81,8 +84,10 @@ def run(self) -> None:
         """Runs the step."""
         # Get the wheel package and add it to the sys path
         wheel_package = self.entrypoint_args[WHEEL_PACKAGE_OPTION]
-        distribution = pkg_resources.get_distribution(wheel_package)
-        project_root = os.path.join(distribution.location, wheel_package)
+
+        dist = distribution(wheel_package)
+        project_root = os.path.join(dist.locate_file("."), wheel_package)
+
         if project_root not in sys.path:
             sys.path.insert(0, project_root)
             sys.path.insert(-1, project_root)

src/zenml/integrations/integration.py

@@ -13,16 +13,14 @@
 #  permissions and limitations under the License.
 """Base and meta classes for ZenML integrations."""
 
-import re
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
 
-import pkg_resources
-from pkg_resources import Requirement
+from packaging.requirements import Requirement
 
 from zenml.integrations.registry import integration_registry
 from zenml.logger import get_logger
 from zenml.stack.flavor import Flavor
-from zenml.utils.integration_utils import parse_requirement
+from zenml.utils.package_utils import get_dependencies, requirement_installed
 
 if TYPE_CHECKING:
     from zenml.plugins.base_plugin_flavor import BasePluginFlavor
@@ -69,65 +67,32 @@ def check_installation(cls) -> bool:
         Returns:
             True if all required packages are installed, False otherwise.
         """
-        for r in cls.get_requirements():
-            try:
-                # First check if the base package is installed
-                dist = pkg_resources.get_distribution(r)
-
-                # Next, check if the dependencies (including extras) are
-                # installed
-                deps: List[Requirement] = []
-
-                _, extras = parse_requirement(r)
-                if extras:
-                    extra_list = extras[1:-1].split(",")
-                    for extra in extra_list:
-                        try:
-                            requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
-                        except pkg_resources.UnknownExtra as e:
-                            logger.debug(f"Unknown extra: {str(e)}")
-                            return False
-                        deps.extend(requirements)
-                else:
-                    deps = dist.requires()
-
-                for ri in deps:
-                    try:
-                        # Remove the "extra == ..." part from the requirement string
-                        cleaned_req = re.sub(
-                            r"; extra == \"\w+\"", "", str(ri)
-                        )
-                        pkg_resources.get_distribution(cleaned_req)
-                    except pkg_resources.DistributionNotFound as e:
-                        logger.debug(
-                            f"Unable to find required dependency "
-                            f"'{e.req}' for requirement '{r}' "
-                            f"necessary for integration '{cls.NAME}'."
-                        )
-                        return False
-                    except pkg_resources.VersionConflict as e:
-                        logger.debug(
-                            f"Package version '{e.dist}' does not match "
-                            f"version '{e.req}' required by '{r}' "
-                            f"necessary for integration '{cls.NAME}'."
-                        )
-                        return False
-
-            except pkg_resources.DistributionNotFound as e:
-                logger.debug(
-                    f"Unable to find required package '{e.req}' for "
-                    f"integration {cls.NAME}."
-                )
-                return False
-            except pkg_resources.VersionConflict as e:
+        for requirement in cls.get_requirements():
+            parsed_requirement = Requirement(requirement)
+
+            if not requirement_installed(parsed_requirement):
                 logger.debug(
-                    f"Package version '{e.dist}' does not match version "
-                    f"'{e.req}' necessary for integration {cls.NAME}."
+                    "Requirement '%s' for integration '%s' is not installed "
+                    "or installed with the wrong version.",
+                    requirement,
+                    cls.NAME,
                 )
                 return False
 
+            dependencies = get_dependencies(parsed_requirement)
+
+            for dependency in dependencies:
+                if not requirement_installed(dependency):
+                    logger.debug(
+                        "Requirement '%s' for integration '%s' is not "
+                        "installed or installed with the wrong version.",
+                        dependency,
+                        cls.NAME,
+                    )
+                    return False
+
         logger.debug(
-            f"Integration {cls.NAME} is installed correctly with "
+            f"Integration '{cls.NAME}' is installed correctly with "
             f"requirements {cls.get_requirements()}."
         )
         return True