Repository: apache/airflow
File: task-sdk/src/airflow/sdk/types.py (Python)

Code before:
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet

if TYPE_CHECKING:
    from collections.abc import Iterator

    from pydantic import AwareDatetime, JsonValue

    from airflow.sdk._shared.logging.types import Logger as Logger
    from airflow.sdk.api.datamodels._generated import PreviousTIResponse, TaskInstanceState
    from airflow.sdk.bases.operator import BaseOperator
    from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAliasEvent, AssetRef, BaseAssetUniqueKey
    from airflow.sdk.definitions.context import Context
    from airflow.sdk.definitions.mappedoperator import MappedOperator

    Operator: TypeAlias = BaseOperator | MappedOperator


class WeightRuleProtocol(Protocol):
    """
    Protocol for weight_rule parameter.

    Accepts str (e.g. "downstream"), WeightRule, or custom strategy instances like PriorityWeightStrategy.
    """

    pass


WeightRuleParam: TypeAlias = str | WeightRule | WeightRuleProtocol


class TaskInstanceKey(NamedTuple):
    """Key used to identify task instance."""

    dag_id: str
    task_id: str
    run_id: str
    try_number: int = 1
    map_index: int = -1

    @property
    def primary(self) -> tuple[str, str, str, int]:
        """Return task instance primary key part of the key."""
        return self.dag_id, self.task_id, self.run_id, self.map_index

    def with_try_number(self, try_number: int) -> TaskInstanceKey:
        """Return TaskInstanceKey with provided ``try_number``."""
        return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, try_number, self.map_index)

    @property
    def key(self) -> TaskInstanceKey:

Review comment (suggestion):
Wouldn’t this simply accept _everything_…? I feel I’m missing something. Additional declarations should be needed to make this protocol to be effective, such as a function signature or an attribute on the protocol.

Code after:
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet

if TYPE_CHECKING:
    from collections.abc import Iterator

    from pydantic import AwareDatetime, JsonValue

    from airflow.models.taskinstance import TaskInstance as SchedulerTaskInstance
    from airflow.sdk._shared.logging.types import Logger as Logger
    from airflow.sdk.api.datamodels._generated import PreviousTIResponse, TaskInstanceState
    from airflow.sdk.bases.operator import BaseOperator
    from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAliasEvent, AssetRef, BaseAssetUniqueKey
    from airflow.sdk.definitions.context import Context
    from airflow.sdk.definitions.mappedoperator import MappedOperator

    Operator: TypeAlias = BaseOperator | MappedOperator


class WeightRuleProtocol(Protocol):
    """
    Protocol for custom weight strategy instances.

    Matches objects that implement get_weight(ti).
    """

    def get_weight(self, ti: SchedulerTaskInstance) -> int:
        """Return the priority weight for the task instance."""
        ...


WeightRuleParam: TypeAlias = str | WeightRule | WeightRuleProtocol


class TaskInstanceKey(NamedTuple):
    """Key used to identify task instance."""

    dag_id: str
    task_id: str
    run_id: str
    try_number: int = 1
    map_index: int = -1

    @property
    def primary(self) -> tuple[str, str, str, int]:
        """Return task instance primary key part of the key."""
        return self.dag_id, self.task_id, self.run_id, self.map_index

    def with_try_number(self, try_number: int) -> TaskInstanceKey:
        """Return TaskInstanceKey with provided ``try_number``."""
        return TaskInstanceKey(self.dag_id, self.task_id, self.run_id, try_number, self.map_index)