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"Version {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"Version {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 |
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"Version {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"Version {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)
]