modify_partners.py

"""A modification workflow that modifies a partner."""

from orchestrator.forms import FormPage
from orchestrator.targets import Target
from orchestrator.types import FormGenerator, State, UUIDstr
from orchestrator.workflow import StepList, begin, done, step, workflow
from pydantic import ConfigDict, EmailStr, field_validator
from pydantic_forms.validators import Choice

from gso.services.partners import (
    ModifiedPartnerSchema,
    edit_partner,
    filter_partners_by_email,
    filter_partners_by_name,
    get_all_partners,
    get_partner_by_name,
)


def initial_input_form_generator() -> FormGenerator:
    """Gather input from the user needed for modifying a partner."""
    partners = {}
    for partner in get_all_partners():
        partners[partner["partner_id"]] = partner["name"]

    partner_choice = Choice("Select a partner", zip(partners.values(), partners.items(), strict=True))  # type: ignore[arg-type]

    class SelectPartnerForm(FormPage):
        model_config = ConfigDict(title="Choose a Partner")

        partners: partner_choice  # type: ignore[valid-type]

    initial_user_input = yield SelectPartnerForm

    partner = get_partner_by_name(name=initial_user_input.partners.name)

    class ModifyPartnerForm(FormPage):
        model_config = ConfigDict(title="Modify a Partner")

        name: str | None = partner["name"]
        email: EmailStr | None = partner["email"]

        @field_validator("name")
        def validate_name(cls, name: str) -> str | None:
            if partner["name"] == name:
                return None
            if filter_partners_by_name(name=name, case_sensitive=False):
                msg = "Partner with this name already exists."
                raise ValueError(msg)
            return name

        @field_validator("email")
        def validate_email(cls, email: str) -> EmailStr | None:
            if partner["email"] == email:
                return None
            if filter_partners_by_email(email=email, case_sensitive=False):
                msg = "Partner with this email already exists."
                raise ValueError(msg)
            return email

    user_input = yield ModifyPartnerForm

    return user_input.model_dump() | {"partner_id": partner["partner_id"]}


@step("Save partner information to database")
def save_partner_to_database(
    partner_id: UUIDstr,
    name: str | None,
    email: EmailStr | None,
) -> State:
    """Save modified partner in database."""
    partner = edit_partner(
        partner_data=ModifiedPartnerSchema(
            partner_id=partner_id,
            name=name,
            email=email,
        )
    )

    return {"modified_partner": partner}


@workflow(
    "Modify partners",
    initial_input_form=initial_input_form_generator,
    target=Target.SYSTEM,
)
def task_modify_partners() -> StepList:
    """Modify a Partner."""
    return begin >> save_partner_to_database >> done