Source code for orchestrator.utils.module_factory

from inspect import isclass
from abc import ABC, abstractmethod
from .recorder import Recorder
from ..utils.exceptions import ModuleAlreadyInFactoryError




class ModuleFactory:
    """
    Factory for spawning concrete modules given a token.

    This class is not instantiated directly, but through a concrete
    :class:`ModuleBuilder`. A factory object is used for one module type,
    specified at instantiation, to avoid token confusion/conflicts. Specific
    module factories, used by module builders, are defined in each module.

    :param base_class: The ABC base class defining the module that this
        factory will spawn.
    :type base_class: ``ABC``
    """



    def __init__(self, base_class):
        self.base_class = base_class
        self._creators = {}




    def add_new_module(self, module_type, module):
        """
        Add a concrete class to the factory.

        :param module_type: Token to specify the module. The convention is to
            use all caps.
        :type module_type: str
        :param module: Implementation of the base_class to be added to the
            factory.
        :type module: type(self.base_class)
        :raises ValueError: If the module_type is already registered or the
            module is not a subclass of the base_class.
        """
        if module_type in self._creators:
            raise ModuleAlreadyInFactoryError(
                f'Trying to add {module_type} to ModuleFactory, but it is '
                'already listed!')
        else:
            if isclass(module) and self.base_class in module.__mro__:
                self._creators[module_type] = module
            else:
                raise ValueError(f'Supplied module is not a '
                                 f'{self.base_class.__name__} class!')




    def select_module(self, module_type):
        """
        Select the type of module to be trained.

        The specified module must have been added to the factory. Available
        modules can be listed with the :meth:`list_modules` function.

        :param module_type: Token to specify the desired module.
        :type module_type: str
        :returns: Requested concrete implementation of a module.
        :rtype: Module
        :raises ValueError: If the module_type has not been added to the
            factory.
        """
        if module_type in self._creators:
            return self._creators[module_type]
        raise ValueError(f'Module type {module_type} has not been added '
                         'to the factory yet! If you expected this module '
                         'to be present, it may be an optional dependency '
                         'which has not been installed in the environment.')




    def list_modules(self):
        """
        Print the currently available modules in the factory.

        This method is primarily designed as a utility for interactive use
        of the orchestrator.
        """
        print('The factory has the following Modules available (Key-Class):')
        for module_type, module_class in self._creators.items():
            print(f'{module_type} - {module_class}')






class ModuleBuilder(Recorder, ABC):
    """
    Abstract base class for module constructors.

    This class sets the factory to be used for the builder. The default is to
    use the specific module_factory generated by specific modules. A user-
    defined :class:`ModuleFactory` can optionally be supplied instead.

    :param factory: A module factory, typically defined in specific
        implementations.
    :type factory: ModuleFactory
    """



    def __init__(self, factory: ModuleFactory):
        """
        Constructor for the ModuleBuilder, sets the factory to build from.

        :param factory: A module factory, typically defined in specific
            implementations.
        :type factory: ModuleFactory
        """
        super().__init__()
        self.factory = factory




    def list_modules(self):
        """
        Print the currently available modules in the builder.

        This method is primarily designed as a utility for interactive use
        of the orchestrator.
        """
        self.factory.list_modules()




    @abstractmethod
    def build(self):
        """
        Return an instance of the specified module.

        The build method signature is not specified by this class, as different
        modules require different parameters to instantiate.
        """
        pass