# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from airflow.sdk.api.datamodels._generated import (
    DagRunState as DagRunState,
    TaskInstanceState as TaskInstanceState,
    TriggerRule as TriggerRule,
    WeightRule as WeightRule,
)
from airflow.sdk.bases.branch import (
    BaseBranchOperator as BaseBranchOperator,
    BranchMixIn as BranchMixIn,
)
from airflow.sdk.bases.hook import BaseHook as BaseHook
from airflow.sdk.bases.notifier import BaseNotifier as BaseNotifier
from airflow.sdk.bases.operator import (
    BaseAsyncOperator as BaseAsyncOperator,
    BaseOperator as BaseOperator,
    chain as chain,
    chain_linear as chain_linear,
    cross_downstream as cross_downstream,
)
from airflow.sdk.bases.operatorlink import BaseOperatorLink as BaseOperatorLink
from airflow.sdk.bases.sensor import (
    BaseSensorOperator as BaseSensorOperator,
    PokeReturnValue as PokeReturnValue,
)
from airflow.sdk.bases.skipmixin import SkipMixin as SkipMixin
from airflow.sdk.bases.xcom import BaseXCom as BaseXCom
from airflow.sdk.configuration import AirflowSDKConfigParser
from airflow.sdk.definitions.asset import (
    Asset as Asset,
    AssetAlias as AssetAlias,
    AssetAll as AssetAll,
    AssetAny as AssetAny,
    AssetWatcher as AssetWatcher,
)
from airflow.sdk.definitions.asset.decorators import asset as asset
from airflow.sdk.definitions.asset.metadata import Metadata as Metadata
from airflow.sdk.definitions.connection import Connection as Connection
from airflow.sdk.definitions.context import (
    Context as Context,
    get_current_context as get_current_context,
    get_parsing_context as get_parsing_context,
)
from airflow.sdk.definitions.dag import DAG as DAG, dag as dag
from airflow.sdk.definitions.decorators import setup as setup, task as task, teardown as teardown
from airflow.sdk.definitions.decorators.task_group import task_group as task_group
from airflow.sdk.definitions.edges import EdgeModifier as EdgeModifier, Label as Label
from airflow.sdk.definitions.param import Param as Param
from airflow.sdk.definitions.partition_mappers.allowed_key import AllowedKeyMapper
from airflow.sdk.definitions.partition_mappers.base import PartitionMapper
from airflow.sdk.definitions.partition_mappers.chain import ChainMapper
from airflow.sdk.definitions.partition_mappers.identity import IdentityMapper
from airflow.sdk.definitions.partition_mappers.product import ProductMapper
from airflow.sdk.definitions.partition_mappers.temporal import (
    StartOfDayMapper,
    StartOfHourMapper,
    StartOfMonthMapper,
    StartOfQuarterMapper,
    StartOfWeekMapper,
    StartOfYearMapper,
)
from airflow.sdk.definitions.retry_policy import (
    ExceptionRetryPolicy as ExceptionRetryPolicy,
    RetryAction as RetryAction,
    RetryDecision as RetryDecision,
    RetryPolicy as RetryPolicy,
    RetryRule as RetryRule,
)
from airflow.sdk.definitions.taskgroup import TaskGroup as TaskGroup
from airflow.sdk.definitions.template import literal as literal
from airflow.sdk.definitions.timetables.assets import (
    AssetOrTimeSchedule,
    PartitionedAssetTimetable,
)
from airflow.sdk.definitions.timetables.events import EventsTimetable
from airflow.sdk.definitions.timetables.interval import (
    CronDataIntervalTimetable,
    DeltaDataIntervalTimetable,
)
from airflow.sdk.definitions.timetables.trigger import (
    CronPartitionTimetable,
    CronTriggerTimetable,
    DeltaTriggerTimetable,
    MultipleCronTriggerTimetable,
)
from airflow.sdk.definitions.variable import Variable as Variable
from airflow.sdk.definitions.xcom_arg import XComArg as XComArg
from airflow.sdk.execution_time import macros as macros
from airflow.sdk.execution_time.cache import SecretCache as SecretCache
from airflow.sdk.io.path import ObjectStoragePath as ObjectStoragePath

conf: AirflowSDKConfigParser

__all__ = [
    "__version__",
    "AllowedKeyMapper",
    "Asset",
    "AssetAlias",
    "AssetAll",
    "AssetAny",
    "AssetOrTimeSchedule",
    "AssetWatcher",
    "BaseAsyncOperator",
    "BaseBranchOperator",
    "BaseHook",
    "BaseNotifier",
    "BaseOperator",
    "BaseOperatorLink",
    "BaseSensorOperator",
    "BaseXCom",
    "BranchMixIn",
    "ChainMapper",
    "Connection",
    "Context",
    "CronDataIntervalTimetable",
    "CronTriggerTimetable",
    "CronPartitionTimetable",
    "DAG",
    "DagRunState",
    "DeltaDataIntervalTimetable",
    "DeltaTriggerTimetable",
    "EdgeModifier",
    "EventsTimetable",
    "ExceptionRetryPolicy",
    "IdentityMapper",
    "Label",
    "Metadata",
    "MultipleCronTriggerTimetable",
    "ObjectStoragePath",
    "Param",
    "PokeReturnValue",
    "PartitionedAssetTimetable",
    "PartitionMapper",
    "ProductMapper",
    "RetryAction",
    "RetryDecision",
    "RetryPolicy",
    "RetryRule",
    "SecretCache",
    "SkipMixin",
    "StartOfDayMapper",
    "StartOfHourMapper",
    "StartOfMonthMapper",
    "StartOfQuarterMapper",
    "StartOfWeekMapper",
    "StartOfYearMapper",
    "TaskGroup",
    "TaskInstanceState",
    "TriggerRule",
    "Variable",
    "WeightRule",
    "XComArg",
    "asset",
    "chain",
    "chain_linear",
    "cross_downstream",
    "dag",
    "get_current_context",
    "get_parsing_context",
    "literal",
    "macros",
    "setup",
    "task",
    "task_group",
    "teardown",
    "conf",
]

__version__: str
