from typing import Iterable, Tuple, Type, Union
from returns.result import ResultE, Success, Failure
from nrel.hive.dispatcher.instruction_generator.instruction_generator import InstructionGenerator
from nrel.hive.model.entity import Entity
from nrel.hive.runner.runner_payload import RunnerPayload
from nrel.hive.state.simulation_state.simulation_state_ops import (
modify_entities_safe as _modify_entities_safe,
)
from nrel.hive.util.fp import throw_or_return
[docs]def update_instruction_generator_safe(
rp: RunnerPayload, ig: InstructionGenerator
) -> ResultE[RunnerPayload]:
"""
Inject an updated InstructionGenerator into a runner payload
Safe method
"""
new_step_fn_or_error = rp.u.step_update.update_instruction_generator(ig)
if isinstance(new_step_fn_or_error, Failure):
return new_step_fn_or_error
else:
new_update = rp.u._replace(step_update=new_step_fn_or_error.unwrap())
new_rp = rp._replace(u=new_update)
return Success(new_rp)
[docs]def update_instruction_generator(rp: RunnerPayload, ig: InstructionGenerator) -> RunnerPayload:
"""
Inject an updated InstructionGenerator into a runner payload
Unsafe method
"""
rp_or_error = update_instruction_generator_safe(rp, ig)
return throw_or_return(rp_or_error)
[docs]def get_instruction_generator_safe(
rp: RunnerPayload, ig: Union[str, Type[InstructionGenerator]]
) -> ResultE[InstructionGenerator]:
"""
Extract an instruction generator from a runner payload
"""
return rp.u.step_update.get_instruction_generator(ig)
[docs]def get_instruction_generator(
rp: RunnerPayload, ig: Union[str, Type[InstructionGenerator]]
) -> InstructionGenerator:
"""
Extract an instruction generator from a runner payload
"""
new_ig_or_error = rp.u.step_update.get_instruction_generator(ig)
return throw_or_return(new_ig_or_error)
[docs]def set_instruction_generators(
rp: RunnerPayload, instruction_generators: Tuple[InstructionGenerator]
) -> RunnerPayload:
"""
Set the instruction generators on a runner payload.
Overwrites any existing instruction generators.
:param rp: The runner payload to update
:param instruction_generators: The instruction generators to set
:return: The updated runner payload
"""
new_step_simulation = rp.u.step_update.update_instruction_generators(instruction_generators)
new_update = rp.u._replace(step_update=new_step_simulation)
new_rp = rp._replace(u=new_update)
return new_rp
[docs]def modify_entities_safe(rp: RunnerPayload, entities: Iterable[Entity]) -> ResultE[RunnerPayload]:
"""
Modify entities in a runner payload
"""
new_s_or_error = _modify_entities_safe(rp.s, entities)
if isinstance(new_s_or_error, Failure):
return new_s_or_error
else:
new_rp = rp._replace(s=new_s_or_error.unwrap())
return Success(new_rp)
[docs]def modify_entities(rp: RunnerPayload, entities: Iterable[Entity]) -> RunnerPayload:
"""
Modify entities in a runner payload
"""
return throw_or_return(modify_entities_safe(rp, entities))