<filename>composapy/dataflow/models.py
from __future__ import annotations
from typing import Optional, Tuple, Dict

from composapy.dataflow.const import ExternalInput
from composapy.dataflow.web import upload_file_to_runs_dir
from composapy.mixins import (
    ObjectSetMixin,
    SessionObjectMixin,
    PandasMixin,
    FileReferenceMixin,
)

from CompAnalytics import Contracts


class ModuleMemberBase(FileReferenceMixin, PandasMixin, SessionObjectMixin):
    """Used as a base class for Input and Result."""

    contract = None  # ModuleInput | ModuleOutput    => union typing issues, leave as is

    def __init__(self, contract, *args, **kwargs):
        self.contract = contract
        super().__init__(*args, **kwargs)

    def _repr_html_(self):
        """Used to display table contracts as pandas dataframes inside of notebooks."""
        if isinstance(self.contract.ValueObj, Contracts.Tables.Table):
            return self.convert_table_to_dataframe(self.contract.ValueObj)._repr_html_()

    @property
    def _type(self) -> any:
        """Returns the contract member, ArgumentType."""
        return self.contract.ArgumentType

    @property
    def type(self) -> str:
        """Returns the contract member, ArgumentTypeFriendlyName."""
        return self.contract.ArgumentTypeFriendlyName

    @property
    def value(self) -> any:
        """Returns the contract member, ValueObj."""
        return self.contract.ValueObj

    @value.setter
    def value(self, val):
        self.contract.ValueObj = val

    @property
    def name(self) -> str:
        """Returns the contract member, Name."""
        return self.contract.Name


class Input(ModuleMemberBase):
    """Wraps ModuleInput contract for a simplified textual user interface."""

    contract: Contracts.ModuleInput

    def __repr__(self):
        repr_value = self.value
        if isinstance(self.value, str):
            repr_value = f"'{self.value}'"
        return f"Input(name={self.name}, type={self.type}, value={repr_value})"


class Result(ModuleMemberBase):
    """Wraps ModuleOutput contract for a simplified textual user interface."""

    contract: Contracts.ModuleOutput

    def __repr__(self):
        repr_value = self.value
        if isinstance(self.value, str):
            repr_value = f"'{self.value}'"
        return f"Result(name='{self.name}', type={self.type}, value={repr_value})"


class InputSet(ObjectSetMixin):
    """Wrapper for objects with parent Module, for convenience methods on the set of items
    contained with self._target.
    """

    _target: Tuple[Input]

    def __init__(self, inputs: Tuple[Input]):
        self._target = inputs


class ResultSet(ObjectSetMixin):
    """Wrapper for objects with parent Module, for convenience methods on the set of items
    contained with self._target.
    """

    _target: Tuple[Result]

    def __init__(self, results: Tuple[Result]):
        self._target = results


class ModuleResultException(Exception):
    pass


class Module(SessionObjectMixin):
    """The object representation of a module inside a dataflow object."""

    contract: Contracts.Module

    def __init__(self, contract: Contracts.Module, *args, **kwargs):
        self.contract = contract
        super().__init__(*args, **kwargs)

    def __repr__(self):
        return f"Module(name='{self.name}', type={self.type})"

    @property
    def type(self) -> str:
        """Returns the contract member, ModuleType.Name."""
        return self.contract.ModuleType.Name

    @property
    def name(self) -> str:
        """Returns the contract member, Name."""
        return self.contract.Name

    @property
    def inputs(self) -> InputSet:  # Dict[str, Input]:
        """Maps each module input, by name, to a corresponding Input object."""
        return InputSet(
            tuple(
                Input(self.contract.ModuleInputs[name], session=self.session)
                for name in self.contract.ModuleInputs.Indexes.Keys
            )
        )

    @property
    def input(self) -> any:
        """Convenience property that gets the first input. Cannot be used if there is more than
        one input.
        """
        if len(self.inputs) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular input, multiple inputs exist. "
                "For modules that contain multiple inputs, please use "
                "inputs instead of input."
            )
        return next(iter(self.inputs))

    @property
    def results(self) -> ResultSet:  # Dict[str, Result]:
        """Maps each module result, by name, to a corresponding Result object."""
        return ResultSet(
            tuple(
                Result(self.contract.ModuleOutputs[name], session=self.session)
                for name in self.contract.ModuleOutputs.Indexes.Keys
            )
        )

    @property
    def result(self) -> any:
        """Convenience property that gets the first result. Cannot be used if there is more than
        one result.
        """
        if len(self.results) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular result, multiple results exist. "
                "For modules that contain multiple results, please use "
                "results instead of result."
            )
        return next(iter(self.results))


class ModuleSet(ObjectSetMixin):
    """Wrapper for objects with parent Module, for convenience methods on the set of items
    contained with self._target.
    """

    _target: Tuple[Module]

    def __init__(self, modules: Tuple[Module]) -> None:
        self._target = modules


class DataFlowRun(SessionObjectMixin):
    """Similar to a DataFlowObject, with a couple of differences. The first difference is that
    every DataFlowRun has an id, where as a DataFlowObject only has an ID if it is saved. The second
    difference is that the modules property on a DataFlowRun returns ModuleSet<ResultModule>
    instead of ModuleSet<Module>, which has the additional functionality of viewing module
    results.
    """

    contract: Contracts.ExecutionState

    def __init__(
        self, execution_state: Contracts.ExecutionState, *args, **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.contract = execution_state

    def __repr__(self):
        return f"DataFlowRun(id={self.id})"

    @property
    def id(self) -> int:
        """Returns the id of dataflow run. Every DataFlowRun is guaranteed to have an id with a
        non-null value.
        """
        return self.contract.Handle.Id

    @property
    def app_id(self) -> int:
        """Returns the originating dataflow application id, assuming originating dataflow was
        saved.
        """
        return self.contract.Handle.AppId

    @property
    def modules(self) -> ModuleSet:
        """A ModuleSet made up of ResultModule's."""
        return ModuleSet(
            tuple(
                Module(_module, session=self.session)
                for _module in self.contract.Application.Modules
            )
        )

    @property
    def module(self) -> any:
        """Convenience property that gets the first module. Cannot be used if there is more than
        one module.
        """
        if len(self.modules) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular module, multiple modules exist. "
                "For dataflow that contain multiple modules, please use "
                "modules instead of module."
            )
        return next(iter(self.modules))


class DataFlowRunSet(ObjectSetMixin):
    """Wrapper for dataflow run objects, using ObjectSetMixin convenience mixin utilities."""

    _target: Tuple[DataFlowRun]

    def __init__(self, dataflow_runs: Tuple[DataFlowRun]) -> None:
        self._target = dataflow_runs


class DataFlowObject(SessionObjectMixin):
    """DataFlowObject can be used to both model and save dataflow configurations, both saved and
    before saving. Holds a reference to the service needed to carry out operations on it's behalf.
    """

    contract: Contracts.Application

    def __init__(self, contract: Contracts.Application, *args, **kwargs):
        self.contract = contract
        super().__init__(*args, **kwargs)

    def __repr__(self):
        return f"DataFlowObject(id={self.id})"

    @property
    def id(self) -> Optional[int]:
        """The contract id. An unsaved DataFlowObject's id property is None."""
        return self.contract.Id

    @property
    def modules(self) -> ModuleSet:
        """A ModuleSet made up of Module's."""
        return ModuleSet(
            tuple(
                Module(_module, session=self.session)
                for _module in self.contract.Modules
            )
        )

    @property
    def module(self) -> any:
        """Convenience property that gets the first module. Cannot be used if there is more than
        one module.
        """
        if len(self.modules) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular module, multiple modules exist. "
                "For dataflow that contain multiple modules, please use "
                "modules instead of module."
            )
        return next(iter(self.modules))

    def save(self) -> DataFlowObject:
        """Saves the contract representation of DataFlowObject, uses server response as the newly
        updated contract object (for instance, saving an unsaved contract will give it an id).
        """
        self.contract: Contracts.Application = self.session.app_service.SaveApplication(
            self.contract
        )
        return self

    def run(self, external_inputs: Dict[str, any] = None) -> DataFlowRun:
        """Runs the dataflow represented by contained contract. Any external modules
        (external int, table, file) that require outside input to run can be added using a
        dictionary with the module input's name and corresponding contract.
        """
        for module in self.modules:
            module.contract.RequestingExecution = True

        self._write_external_inputs(external_inputs)

        execution_handle: Contracts.ExecutionHandle = (
            self.session.app_service.CreateExecutionContext(
                self.contract, Contracts.ExecutionContextOptions()
            )
        )

        self._post_context_setup_steps(external_inputs, execution_handle)

        self.session.app_service.RunExecutionContext(execution_handle)
        post_run_execution_state: Contracts.ExecutionState = (
            self.session.app_service.GetRun(execution_handle.Id)
        )

        dataflow_run = DataFlowRun(post_run_execution_state, session=self.session)
        return dataflow_run

    def _write_external_inputs(self, external_inputs: dict):
        """Update any module inputs with required external inputs."""
        if not external_inputs:
            return

        for module in self.modules:
            if module.type != ExternalInput.FILE and module.type in ExternalInput.ALL:
                input_name = module.contract.ModuleInputs["Name"].ValueObj
                if input_name in external_inputs.keys():
                    module.contract.ModuleInputs["Input"].ValueObj = external_inputs[
                        input_name
                    ]

    def _post_context_setup_steps(
        self, external_inputs: dict, execution_handle: Contracts.ExecutionHandle
    ):
        """Updates necessary post-context setup item for each individual module, such as
        uploading any needed files before running the execution context."""
        if not external_inputs:
            return

        for module in self.modules:
            if module.contract.ModuleType.Name == ExternalInput.FILE:
                upload_file_to_runs_dir(
                    self.session, execution_handle, module, external_inputs
                )

<filename>composapy/dataflow/models.py
from __future__ import annotations
from typing import Optional, Tuple, Dict

from composapy.dataflow.const import ExternalInput
from composapy.dataflow.web import upload_file_to_runs_dir
from composapy.mixins import (
    ObjectSetMixin,
    SessionObjectMixin,
    PandasMixin,
    FileReferenceMixin,
)

from CompAnalytics import Contracts


class ModuleMemberBase(FileReferenceMixin, PandasMixin, SessionObjectMixin):
    """Used as a base class for Input and Result."""

    contract = None  # ModuleInput | ModuleOutput    => union typing issues, leave as is

    def __init__(self, contract, *args, **kwargs):
        self.contract = contract
        super().__init__(*args, **kwargs)

    def _repr_html_(self):
        """Used to display table contracts as pandas dataframes inside of notebooks."""
        if isinstance(self.contract.ValueObj, Contracts.Tables.Table):
            return self.convert_table_to_dataframe(self.contract.ValueObj)._repr_html_()

    @property
    def _type(self) -> any:
        """Returns the contract member, ArgumentType."""
        return self.contract.ArgumentType

    @property
    def type(self) -> str:
        """Returns the contract member, ArgumentTypeFriendlyName."""
        return self.contract.ArgumentTypeFriendlyName

    @property
    def value(self) -> any:
        """Returns the contract member, ValueObj."""
        return self.contract.ValueObj

    @value.setter
    def value(self, val):
        self.contract.ValueObj = val

    @property
    def name(self) -> str:
        """Returns the contract member, Name."""
        return self.contract.Name


class Input(ModuleMemberBase):
    """Wraps ModuleInput contract for a simplified textual user interface."""

    contract: Contracts.ModuleInput

    def __repr__(self):
        repr_value = self.value
        if isinstance(self.value, str):
            repr_value = f"'{self.value}'"
        return f"Input(name={self.name}, type={self.type}, value={repr_value})"


class Result(ModuleMemberBase):
    """Wraps ModuleOutput contract for a simplified textual user interface."""

    contract: Contracts.ModuleOutput

    def __repr__(self):
        repr_value = self.value
        if isinstance(self.value, str):
            repr_value = f"'{self.value}'"
        return f"Result(name='{self.name}', type={self.type}, value={repr_value})"


class InputSet(ObjectSetMixin):
    """Wrapper for objects with parent Module, for convenience methods on the set of items
    contained with self._target.
    """

    _target: Tuple[Input]

    def __init__(self, inputs: Tuple[Input]):
        self._target = inputs


class ResultSet(ObjectSetMixin):
    """Wrapper for objects with parent Module, for convenience methods on the set of items
    contained with self._target.
    """

    _target: Tuple[Result]

    def __init__(self, results: Tuple[Result]):
        self._target = results


class ModuleResultException(Exception):
    pass


class Module(SessionObjectMixin):
    """The object representation of a module inside a dataflow object."""

    contract: Contracts.Module

    def __init__(self, contract: Contracts.Module, *args, **kwargs):
        self.contract = contract
        super().__init__(*args, **kwargs)

    def __repr__(self):
        return f"Module(name='{self.name}', type={self.type})"

    @property
    def type(self) -> str:
        """Returns the contract member, ModuleType.Name."""
        return self.contract.ModuleType.Name

    @property
    def name(self) -> str:
        """Returns the contract member, Name."""
        return self.contract.Name

    @property
    def inputs(self) -> InputSet:  # Dict[str, Input]:
        """Maps each module input, by name, to a corresponding Input object."""
        return InputSet(
            tuple(
                Input(self.contract.ModuleInputs[name], session=self.session)
                for name in self.contract.ModuleInputs.Indexes.Keys
            )
        )

    @property
    def input(self) -> any:
        """Convenience property that gets the first input. Cannot be used if there is more than
        one input.
        """
        if len(self.inputs) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular input, multiple inputs exist. "
                "For modules that contain multiple inputs, please use "
                "inputs instead of input."
            )
        return next(iter(self.inputs))

    @property
    def results(self) -> ResultSet:  # Dict[str, Result]:
        """Maps each module result, by name, to a corresponding Result object."""
        return ResultSet(
            tuple(
                Result(self.contract.ModuleOutputs[name], session=self.session)
                for name in self.contract.ModuleOutputs.Indexes.Keys
            )
        )

    @property
    def result(self) -> any:
        """Convenience property that gets the first result. Cannot be used if there is more than
        one result.
        """
        if len(self.results) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular result, multiple results exist. "
                "For modules that contain multiple results, please use "
                "results instead of result."
            )
        return next(iter(self.results))


class ModuleSet(ObjectSetMixin):
    """Wrapper for objects with parent Module, for convenience methods on the set of items
    contained with self._target.
    """

    _target: Tuple[Module]

    def __init__(self, modules: Tuple[Module]) -> None:
        self._target = modules


class DataFlowRun(SessionObjectMixin):
    """Similar to a DataFlowObject, with a couple of differences. The first difference is that
    every DataFlowRun has an id, where as a DataFlowObject only has an ID if it is saved. The second
    difference is that the modules property on a DataFlowRun returns ModuleSet<ResultModule>
    instead of ModuleSet<Module>, which has the additional functionality of viewing module
    results.
    """

    contract: Contracts.ExecutionState

    def __init__(
        self, execution_state: Contracts.ExecutionState, *args, **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.contract = execution_state

    def __repr__(self):
        return f"DataFlowRun(id={self.id})"

    @property
    def id(self) -> int:
        """Returns the id of dataflow run. Every DataFlowRun is guaranteed to have an id with a
        non-null value.
        """
        return self.contract.Handle.Id

    @property
    def app_id(self) -> int:
        """Returns the originating dataflow application id, assuming originating dataflow was
        saved.
        """
        return self.contract.Handle.AppId

    @property
    def modules(self) -> ModuleSet:
        """A ModuleSet made up of ResultModule's."""
        return ModuleSet(
            tuple(
                Module(_module, session=self.session)
                for _module in self.contract.Application.Modules
            )
        )

    @property
    def module(self) -> any:
        """Convenience property that gets the first module. Cannot be used if there is more than
        one module.
        """
        if len(self.modules) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular module, multiple modules exist. "
                "For dataflow that contain multiple modules, please use "
                "modules instead of module."
            )
        return next(iter(self.modules))


class DataFlowRunSet(ObjectSetMixin):
    """Wrapper for dataflow run objects, using ObjectSetMixin convenience mixin utilities."""

    _target: Tuple[DataFlowRun]

    def __init__(self, dataflow_runs: Tuple[DataFlowRun]) -> None:
        self._target = dataflow_runs


class DataFlowObject(SessionObjectMixin):
    """DataFlowObject can be used to both model and save dataflow configurations, both saved and
    before saving. Holds a reference to the service needed to carry out operations on it's behalf.
    """

    contract: Contracts.Application

    def __init__(self, contract: Contracts.Application, *args, **kwargs):
        self.contract = contract
        super().__init__(*args, **kwargs)

    def __repr__(self):
        return f"DataFlowObject(id={self.id})"

    @property
    def id(self) -> Optional[int]:
        """The contract id. An unsaved DataFlowObject's id property is None."""
        return self.contract.Id

    @property
    def modules(self) -> ModuleSet:
        """A ModuleSet made up of Module's."""
        return ModuleSet(
            tuple(
                Module(_module, session=self.session)
                for _module in self.contract.Modules
            )
        )

    @property
    def module(self) -> any:
        """Convenience property that gets the first module. Cannot be used if there is more than
        one module.
        """
        if len(self.modules) > 1:
            raise ModuleResultException(
                "Unable to retrieve singular module, multiple modules exist. "
                "For dataflow that contain multiple modules, please use "
                "modules instead of module."
            )
        return next(iter(self.modules))

    def save(self) -> DataFlowObject:
        """Saves the contract representation of DataFlowObject, uses server response as the newly
        updated contract object (for instance, saving an unsaved contract will give it an id).
        """
        self.contract: Contracts.Application = self.session.app_service.SaveApplication(
            self.contract
        )
        return self

    def run(self, external_inputs: Dict[str, any] = None) -> DataFlowRun:
        """Runs the dataflow represented by contained contract. Any external modules
        (external int, table, file) that require outside input to run can be added using a
        dictionary with the module input's name and corresponding contract.
        """
        for module in self.modules:
            module.contract.RequestingExecution = True

        self._write_external_inputs(external_inputs)

        execution_handle: Contracts.ExecutionHandle = (
            self.session.app_service.CreateExecutionContext(
                self.contract, Contracts.ExecutionContextOptions()
            )
        )

        self._post_context_setup_steps(external_inputs, execution_handle)

        self.session.app_service.RunExecutionContext(execution_handle)
        post_run_execution_state: Contracts.ExecutionState = (
            self.session.app_service.GetRun(execution_handle.Id)
        )

        dataflow_run = DataFlowRun(post_run_execution_state, session=self.session)
        return dataflow_run

    def _write_external_inputs(self, external_inputs: dict):
        """Update any module inputs with required external inputs."""
        if not external_inputs:
            return

        for module in self.modules:
            if module.type != ExternalInput.FILE and module.type in ExternalInput.ALL:
                input_name = module.contract.ModuleInputs["Name"].ValueObj
                if input_name in external_inputs.keys():
                    module.contract.ModuleInputs["Input"].ValueObj = external_inputs[
                        input_name
                    ]

    def _post_context_setup_steps(
        self, external_inputs: dict, execution_handle: Contracts.ExecutionHandle
    ):
        """Updates necessary post-context setup item for each individual module, such as
        uploading any needed files before running the execution context."""
        if not external_inputs:
            return

        for module in self.modules:
            if module.contract.ModuleType.Name == ExternalInput.FILE:
                upload_file_to_runs_dir(
                    self.session, execution_handle, module, external_inputs
                )

Used as a base class for Input and Result. # ModuleInput | ModuleOutput    => union typing issues, leave as is Used to display table contracts as pandas dataframes inside of notebooks. Returns the contract member, ArgumentType. Returns the contract member, ArgumentTypeFriendlyName. Returns the contract member, ValueObj. Returns the contract member, Name. Wraps ModuleInput contract for a simplified textual user interface. Wraps ModuleOutput contract for a simplified textual user interface. Wrapper for objects with parent Module, for convenience methods on the set of items     contained with self._target. Wrapper for objects with parent Module, for convenience methods on the set of items     contained with self._target. The object representation of a module inside a dataflow object. Returns the contract member, ModuleType.Name. Returns the contract member, Name. # Dict[str, Input]: Maps each module input, by name, to a corresponding Input object. Convenience property that gets the first input. Cannot be used if there is more than         one input. # Dict[str, Result]: Maps each module result, by name, to a corresponding Result object. Convenience property that gets the first result. Cannot be used if there is more than         one result. Wrapper for objects with parent Module, for convenience methods on the set of items     contained with self._target. Similar to a DataFlowObject, with a couple of differences. The first difference is that     every DataFlowRun has an id, where as a DataFlowObject only has an ID if it is saved. The second     difference is that the modules property on a DataFlowRun returns ModuleSet<ResultModule>     instead of ModuleSet<Module>, which has the additional functionality of viewing module     results. Returns the id of dataflow run. Every DataFlowRun is guaranteed to have an id with a         non-null value. Returns the originating dataflow application id, assuming originating dataflow was         saved. A ModuleSet made up of ResultModule's. Convenience property that gets the first module. Cannot be used if there is more than         one module. Wrapper for dataflow run objects, using ObjectSetMixin convenience mixin utilities. DataFlowObject can be used to both model and save dataflow configurations, both saved and     before saving. Holds a reference to the service needed to carry out operations on it's behalf. The contract id. An unsaved DataFlowObject's id property is None. A ModuleSet made up of Module's. Convenience property that gets the first module. Cannot be used if there is more than         one module. Saves the contract representation of DataFlowObject, uses server response as the newly         updated contract object (for instance, saving an unsaved contract will give it an id). Runs the dataflow represented by contained contract. Any external modules         (external int, table, file) that require outside input to run can be added using a         dictionary with the module input's name and corresponding contract. Update any module inputs with required external inputs. Updates necessary post-context setup item for each individual module, such as         uploading any needed files before running the execution context.