Source code for airflow_pydantic.operators.weekday

from logging import getLogger

from pydantic import Field, field_validator

from ..core import Task, TaskArgs
from ..utils import DatetimeArg, ImportPath

__all__ = (
    "BranchDayOfWeekOperator",
    "BranchDayOfWeekOperatorArgs",
    "BranchDayOfWeekTask",
    "BranchDayOfWeekTaskArgs",
)

_log = getLogger(__name__)




class BranchDayOfWeekTaskArgs(TaskArgs):
    # https://airflow.apache.org/docs/apache-airflow-providers-standard/stable/_api/airflow/providers/standard/operators/weekday/index.html#airflow.providers.standard.operators.weekday.BranchDayOfWeekOperator
    follow_task_ids_if_true: list[str] | None = Field(default=None, description="List of task IDs to follow if condition evaluates to True")
    follow_task_ids_if_false: list[str] | None = Field(default=None, description="List of task IDs to follow if condition evaluates to False")
    target_lower: DatetimeArg | None = Field(default=None, description="The lower bound datetime to compare against")
    target_upper: DatetimeArg | None = Field(default=None, description="The upper bound datetime to compare against")
    use_task_logical_date: bool | None = Field(
        default=None, description="If True, uses the task's logical date for comparison; otherwise, uses the current datetime"
    )
    use_task_execution_date: bool | None = Field(
        default=None, description="If True, uses the task's execution date for comparison; otherwise, uses the current datetime"
    )



# Alias
BranchDayOfWeekOperatorArgs = BranchDayOfWeekTaskArgs




class BranchDayOfWeekTask(Task, BranchDayOfWeekTaskArgs):
    operator: ImportPath = Field(
        default="airflow_pydantic.airflow.BranchDayOfWeekOperator", description="airflow operator path", validate_default=True
    )

    @field_validator("operator")
    @classmethod
    def validate_operator(cls, v: type) -> type:
        from airflow_pydantic.airflow import BranchDayOfWeekOperator, _AirflowPydanticMarker

        if not isinstance(v, type):
            raise TypeError(f"operator must be 'airflow.providers.standard.operators.weekday.BranchDayOfWeekOperator', got: {v}")
        if issubclass(v, _AirflowPydanticMarker):
            _log.info("BranchDateTimeOperator is a marker class, returning as is")
            return v
        if not issubclass(v, BranchDayOfWeekOperator):
            raise TypeError(f"operator must be 'airflow.providers.standard.operators.weekday.BranchDayOfWeekOperator', got: {v}")
        return v



# Alias
BranchDayOfWeekOperator = BranchDayOfWeekTask