Registry

zenml.integrations.registry

Implementation of a registry to track ZenML integrations.

IntegrationRegistry

Registry to keep track of ZenML Integrations.

Source code in zenml/integrations/registry.py

class IntegrationRegistry(object):
    """Registry to keep track of ZenML Integrations."""

    def __init__(self) -> None:
        """Initializing the integration registry."""
        self._integrations: Dict[str, Type["Integration"]] = {}

    @property
    def integrations(self) -> Dict[str, Type["Integration"]]:
        """Method to get integrations dictionary.

        Returns:
            A dict of integration key to type of `Integration`.
        """
        return self._integrations

    @integrations.setter
    def integrations(self, i: Any) -> None:
        """Setter method for the integrations property.

        Args:
            i: Value to set the integrations property to.

        Raises:
            IntegrationError: If you try to manually set the integrations property.
        """
        raise IntegrationError(
            "Please do not manually change the integrations within the "
            "registry. If you would like to register a new integration "
            "manually, please use "
            "`integration_registry.register_integration()`."
        )

    def register_integration(
        self, key: str, type_: Type["Integration"]
    ) -> None:
        """Method to register an integration with a given name.

        Args:
            key: Name of the integration.
            type_: Type of the integration.
        """
        self._integrations[key] = type_

    def activate_integrations(self) -> None:
        """Method to activate the integrations with are registered in the registry."""
        for name, integration in self._integrations.items():
            if integration.check_installation():
                logger.debug(f"Activating integration `{name}`...")
                integration.activate()
                logger.debug(f"Integration `{name}` is activated.")
            else:
                logger.debug(f"Integration `{name}` could not be activated.")

    @property
    def list_integration_names(self) -> List[str]:
        """Get a list of all possible integrations.

        Returns:
            A list of all possible integrations.
        """
        return [name for name in self._integrations]

    def select_integration_requirements(
        self,
        integration_name: Optional[str] = None,
        target_os: Optional[str] = None,
    ) -> List[str]:
        """Select the requirements for a given integration or all integrations.

        Args:
            integration_name: Name of the integration to check.
            target_os: Target OS for the requirements.

        Returns:
            List of requirements for the integration.

        Raises:
            KeyError: If the integration is not found.
        """
        if integration_name:
            if integration_name in self.list_integration_names:
                return self._integrations[integration_name].get_requirements(
                    target_os=target_os
                )
            else:
                raise KeyError(
                    f"Integration {integration_name} does not exist. "
                    f"Currently the following integrations are implemented. "
                    f"{self.list_integration_names}"
                )
        else:
            return [
                requirement
                for name in self.list_integration_names
                for requirement in self._integrations[name].get_requirements(
                    target_os=target_os
                )
            ]

    def select_uninstall_requirements(
        self,
        integration_name: Optional[str] = None,
        target_os: Optional[str] = None,
    ) -> List[str]:
        """Select the uninstall requirements for a given integration or all integrations.

        Args:
            integration_name: Name of the integration to check.
            target_os: Target OS for the requirements.

        Returns:
            List of requirements for the integration uninstall.

        Raises:
            KeyError: If the integration is not found.
        """
        if integration_name:
            if integration_name in self.list_integration_names:
                return self._integrations[
                    integration_name
                ].get_uninstall_requirements(target_os=target_os)
            else:
                raise KeyError(
                    f"Integration {integration_name} does not exist. "
                    f"Currently the following integrations are implemented. "
                    f"{self.list_integration_names}"
                )
        else:
            return [
                requirement
                for name in self.list_integration_names
                for requirement in self._integrations[
                    name
                ].get_uninstall_requirements(target_os=target_os)
            ]

    def is_installed(self, integration_name: Optional[str] = None) -> bool:
        """Checks if all requirements for an integration are installed.

        Args:
            integration_name: Name of the integration to check.

        Returns:
            True if all requirements are installed, False otherwise.

        Raises:
            KeyError: If the integration is not found.
        """
        if integration_name in self.list_integration_names:
            return self._integrations[integration_name].check_installation()
        elif not integration_name:
            all_installed = [
                self._integrations[item].check_installation()
                for item in self.list_integration_names
            ]
            return all(all_installed)
        else:
            raise KeyError(
                f"Integration '{integration_name}' not found. "
                f"Currently the following integrations are available: "
                f"{self.list_integration_names}"
            )

    def get_installed_integrations(self) -> List[str]:
        """Returns list of installed integrations.

        Returns:
            List of installed integrations.
        """
        return [
            name
            for name, integration in integration_registry.integrations.items()
            if integration.check_installation()
        ]

integrations: Dict[str, Type[Integration]] property writable

Method to get integrations dictionary.

Returns:

Type	Description
Dict[str, Type[Integration]]	A dict of integration key to type of Integration.

list_integration_names: List[str] property readonly

Get a list of all possible integrations.

Returns:

Type	Description
List[str]	A list of all possible integrations.

__init__(self) special

Initializing the integration registry.

Source code in zenml/integrations/registry.py

def __init__(self) -> None:
    """Initializing the integration registry."""
    self._integrations: Dict[str, Type["Integration"]] = {}

activate_integrations(self)

Method to activate the integrations with are registered in the registry.

Source code in zenml/integrations/registry.py

def activate_integrations(self) -> None:
    """Method to activate the integrations with are registered in the registry."""
    for name, integration in self._integrations.items():
        if integration.check_installation():
            logger.debug(f"Activating integration `{name}`...")
            integration.activate()
            logger.debug(f"Integration `{name}` is activated.")
        else:
            logger.debug(f"Integration `{name}` could not be activated.")

get_installed_integrations(self)

Returns list of installed integrations.

Returns:

Type	Description
List[str]	List of installed integrations.

Source code in zenml/integrations/registry.py

def get_installed_integrations(self) -> List[str]:
    """Returns list of installed integrations.

    Returns:
        List of installed integrations.
    """
    return [
        name
        for name, integration in integration_registry.integrations.items()
        if integration.check_installation()
    ]

is_installed(self, integration_name=None)

Checks if all requirements for an integration are installed.

Parameters:

Name	Type	Description	Default
integration_name	Optional[str]	Name of the integration to check.	None

Returns:

Type	Description
bool	True if all requirements are installed, False otherwise.

Exceptions:

Type	Description
KeyError	If the integration is not found.

Source code in zenml/integrations/registry.py

def is_installed(self, integration_name: Optional[str] = None) -> bool:
    """Checks if all requirements for an integration are installed.

    Args:
        integration_name: Name of the integration to check.

    Returns:
        True if all requirements are installed, False otherwise.

    Raises:
        KeyError: If the integration is not found.
    """
    if integration_name in self.list_integration_names:
        return self._integrations[integration_name].check_installation()
    elif not integration_name:
        all_installed = [
            self._integrations[item].check_installation()
            for item in self.list_integration_names
        ]
        return all(all_installed)
    else:
        raise KeyError(
            f"Integration '{integration_name}' not found. "
            f"Currently the following integrations are available: "
            f"{self.list_integration_names}"
        )

register_integration(self, key, type_)

Method to register an integration with a given name.

Parameters:

Name	Type	Description	Default
key	str	Name of the integration.	required
type_	Type[Integration]	Type of the integration.	required

Source code in zenml/integrations/registry.py

def register_integration(
    self, key: str, type_: Type["Integration"]
) -> None:
    """Method to register an integration with a given name.

    Args:
        key: Name of the integration.
        type_: Type of the integration.
    """
    self._integrations[key] = type_

select_integration_requirements(self, integration_name=None, target_os=None)

Select the requirements for a given integration or all integrations.

Parameters:

Name	Type	Description	Default
integration_name	Optional[str]	Name of the integration to check.	None
target_os	Optional[str]	Target OS for the requirements.	None

Returns:

Type	Description
List[str]	List of requirements for the integration.

Exceptions:

Type	Description
KeyError	If the integration is not found.

Source code in zenml/integrations/registry.py

def select_integration_requirements(
    self,
    integration_name: Optional[str] = None,
    target_os: Optional[str] = None,
) -> List[str]:
    """Select the requirements for a given integration or all integrations.

    Args:
        integration_name: Name of the integration to check.
        target_os: Target OS for the requirements.

    Returns:
        List of requirements for the integration.

    Raises:
        KeyError: If the integration is not found.
    """
    if integration_name:
        if integration_name in self.list_integration_names:
            return self._integrations[integration_name].get_requirements(
                target_os=target_os
            )
        else:
            raise KeyError(
                f"Integration {integration_name} does not exist. "
                f"Currently the following integrations are implemented. "
                f"{self.list_integration_names}"
            )
    else:
        return [
            requirement
            for name in self.list_integration_names
            for requirement in self._integrations[name].get_requirements(
                target_os=target_os
            )
        ]

select_uninstall_requirements(self, integration_name=None, target_os=None)

Select the uninstall requirements for a given integration or all integrations.

Parameters:

Name	Type	Description	Default
integration_name	Optional[str]	Name of the integration to check.	None
target_os	Optional[str]	Target OS for the requirements.	None

Returns:

Type	Description
List[str]	List of requirements for the integration uninstall.

Exceptions:

Type	Description
KeyError	If the integration is not found.

Source code in zenml/integrations/registry.py

def select_uninstall_requirements(
    self,
    integration_name: Optional[str] = None,
    target_os: Optional[str] = None,
) -> List[str]:
    """Select the uninstall requirements for a given integration or all integrations.

    Args:
        integration_name: Name of the integration to check.
        target_os: Target OS for the requirements.

    Returns:
        List of requirements for the integration uninstall.

    Raises:
        KeyError: If the integration is not found.
    """
    if integration_name:
        if integration_name in self.list_integration_names:
            return self._integrations[
                integration_name
            ].get_uninstall_requirements(target_os=target_os)
        else:
            raise KeyError(
                f"Integration {integration_name} does not exist. "
                f"Currently the following integrations are implemented. "
                f"{self.list_integration_names}"
            )
    else:
        return [
            requirement
            for name in self.list_integration_names
            for requirement in self._integrations[
                name
            ].get_uninstall_requirements(target_os=target_os)
        ]