metaflow-stubs 2.12.28__py2.py3-none-any.whl → 2.12.29__py2.py3-none-any.whl

metaflow-stubs/plugins/airflow/airflow.pyi

@@ -1,36 +1,26 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.504876                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.622880                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
-import typing
-if typing.TYPE_CHECKING:
-    import metaflow._vendor.click.types
-    import metaflow.metaflow_current
-    import metaflow.exception
 
-current: metaflow.metaflow_current.Current
-
-class MetaflowException(Exception, metaclass=type):
-    def __init__(self, msg = "", lineno = None):
-        ...
-    def __str__(self):
-        ...
-    ...
-
-class FilePathClass(metaflow._vendor.click.types.ParamType, metaclass=type):
-    def __init__(self, is_text, encoding):
-        ...
-    def convert(self, value, param, ctx):
-        ...
-    def __str__(self):
-        ...
-    def __repr__(self):
-        ...
-    ...
+from ...metaflow_current import current as current
+from ...exception import MetaflowException as MetaflowException
+from ...includefile import FilePathClass as FilePathClass
+from ...parameters import DelayedEvaluationParameter as DelayedEvaluationParameter
+from ...parameters import JSONTypeClass as JSONTypeClass
+from ...parameters import deploy_time_eval as deploy_time_eval
+from ..cards.card_modules import chevron as chevron
+from ..kubernetes.kubernetes import Kubernetes as Kubernetes
+from ..timeout_decorator import get_run_time_limit_for_task as get_run_time_limit_for_task
+from . import airflow_utils as airflow_utils
+from .airflow_utils import AIRFLOW_MACROS as AIRFLOW_MACROS
+from .airflow_utils import AirflowTask as AirflowTask
+from .airflow_utils import Workflow as Workflow
+from .exception import AirflowException as AirflowException
 
 AIRFLOW_KUBERNETES_CONN_ID: None
 
@@ -74,109 +64,8 @@ SERVICE_INTERNAL_URL: None
 
 AZURE_KEY_VAULT_PREFIX: None
 
-class DelayedEvaluationParameter(object, metaclass=type):
-    """
-    This is a very simple wrapper to allow parameter "conversion" to be delayed until
-    the `_set_constants` function in FlowSpec. Typically, parameters are converted
-    by click when the command line option is processed. For some parameters, like
-    IncludeFile, this is too early as it would mean we would trigger the upload
-    of the file too early. If a parameter converts to a DelayedEvaluationParameter
-    object through the usual click mechanisms, `_set_constants` knows to invoke the
-    __call__ method on that DelayedEvaluationParameter; in that case, the __call__
-    method is invoked without any parameter. The return_str parameter will be used
-    by schedulers when they need to convert DelayedEvaluationParameters to a
-    string to store them
-    """
-    def __init__(self, name, field, fun):
-        ...
-    def __call__(self, return_str = False):
-        ...
-    ...
-
-class JSONTypeClass(metaflow._vendor.click.types.ParamType, metaclass=type):
-    def convert(self, value, param, ctx):
-        ...
-    def __str__(self):
-        ...
-    def __repr__(self):
-        ...
-    ...
-
-def deploy_time_eval(value):
-    ...
-
-class Kubernetes(object, metaclass=type):
-    def __init__(self, datastore, metadata, environment):
-        ...
-    def launch_job(self, **kwargs):
-        ...
-    def create_jobset(self, flow_name, run_id, step_name, task_id, attempt, user, code_package_sha, code_package_url, code_package_ds, docker_image, docker_image_pull_policy, step_cli = None, service_account = None, secrets = None, node_selector = None, namespace = None, cpu = None, gpu = None, gpu_vendor = None, disk = None, memory = None, use_tmpfs = None, tmpfs_tempdir = None, tmpfs_size = None, tmpfs_path = None, run_time_limit = None, env = None, persistent_volume_claims = None, tolerations = None, labels = None, shared_memory = None, port = None, num_parallel = None):
-        ...
-    def create_job_object(self, flow_name, run_id, step_name, task_id, attempt, user, code_package_sha, code_package_url, code_package_ds, step_cli, docker_image, docker_image_pull_policy, service_account = None, secrets = None, node_selector = None, namespace = None, cpu = None, gpu = None, gpu_vendor = None, disk = None, memory = None, use_tmpfs = None, tmpfs_tempdir = None, tmpfs_size = None, tmpfs_path = None, run_time_limit = None, env = None, persistent_volume_claims = None, tolerations = None, labels = None, shared_memory = None, port = None, name_pattern = None):
-        ...
-    def create_k8sjob(self, job):
-        ...
-    def wait(self, stdout_location, stderr_location, echo = None):
-        ...
-    ...
-
-def get_run_time_limit_for_task(step_decos):
-    ...
-
-class AIRFLOW_MACROS(object, metaclass=type):
-    @classmethod
-    def create_task_id(cls, is_foreach):
-        ...
-    @classmethod
-    def pathspec(cls, flowname, is_foreach = False):
-        ...
-    ...
-
 TASK_ID_XCOM_KEY: str
 
-class AirflowTask(object, metaclass=type):
-    def __init__(self, name, operator_type = "kubernetes", flow_name = None, is_mapper_node = False, flow_contains_foreach = False):
-        ...
-    @property
-    def is_mapper_node(self):
-        ...
-    def set_operator_args(self, **kwargs):
-        ...
-    def to_dict(self):
-        ...
-    @classmethod
-    def from_dict(cls, task_dict, flow_name = None, flow_contains_foreach = False):
-        ...
-    def to_task(self):
-        ...
-    ...
-
-class Workflow(object, metaclass=type):
-    def __init__(self, file_path = None, graph_structure = None, metadata = None, **kwargs):
-        ...
-    def set_parameters(self, params):
-        ...
-    def add_state(self, state):
-        ...
-    def to_dict(self):
-        ...
-    def to_json(self):
-        ...
-    @classmethod
-    def from_dict(cls, data_dict):
-        ...
-    @classmethod
-    def from_json(cls, json_string):
-        ...
-    def compile(self):
-        ...
-    ...
-
-class AirflowException(metaflow.exception.MetaflowException, metaclass=type):
-    def __init__(self, msg):
-        ...
-    ...
-
 SUPPORTED_SENSORS: list
 
 AIRFLOW_DEPLOY_TEMPLATE_FILE: str

metaflow-stubs/plugins/airflow/airflow_cli.pyi

@@ -1,146 +1,27 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.506047                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.623352                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
+import metaflow
 import typing
 if typing.TYPE_CHECKING:
-    import metaflow.decorators
-    import metaflow.metaflow_current
     import metaflow.exception
 
-current: metaflow.metaflow_current.Current
-
-class MetaflowException(Exception, metaclass=type):
-    def __init__(self, msg = "", lineno = None):
-        ...
-    def __str__(self):
-        ...
-    ...
-
-class MetaflowInternalError(metaflow.exception.MetaflowException, metaclass=type):
-    ...
-
-def load_token(token_prefix):
-    ...
-
-def new_token(token_prefix, prev_token = None):
-    ...
-
-def store_token(token_prefix, token):
-    ...
-
-class KubernetesDecorator(metaflow.decorators.StepDecorator, metaclass=type):
-    """
-    Specifies that this step should execute on Kubernetes.
-    
-    Parameters
-    ----------
-    cpu : int, default 1
-        Number of CPUs required for this step. If `@resources` is
-        also present, the maximum value from all decorators is used.
-    memory : int, default 4096
-        Memory size (in MB) required for this step. If
-        `@resources` is also present, the maximum value from all decorators is
-        used.
-    disk : int, default 10240
-        Disk size (in MB) required for this step. If
-        `@resources` is also present, the maximum value from all decorators is
-        used.
-    image : str, optional, default None
-        Docker image to use when launching on Kubernetes. If not specified, and
-        METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
-        not, a default Docker image mapping to the current version of Python is used.
-    image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
-        If given, the imagePullPolicy to be applied to the Docker image of the step.
-    service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
-        Kubernetes service account to use when launching pod in Kubernetes.
-    secrets : List[str], optional, default None
-        Kubernetes secrets to use when launching pod in Kubernetes. These
-        secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
-        in Metaflow configuration.
-    node_selector: Union[Dict[str,str], str], optional, default None
-        Kubernetes node selector(s) to apply to the pod running the task.
-        Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
-        or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
-    namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
-        Kubernetes namespace to use when launching pod in Kubernetes.
-    gpu : int, optional, default None
-        Number of GPUs required for this step. A value of zero implies that
-        the scheduled node should not have GPUs.
-    gpu_vendor : str, default KUBERNETES_GPU_VENDOR
-        The vendor of the GPUs to be used for this step.
-    tolerations : List[str], default []
-        The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
-        Kubernetes tolerations to use when launching pod in Kubernetes.
-    use_tmpfs : bool, default False
-        This enables an explicit tmpfs mount for this step.
-    tmpfs_tempdir : bool, default True
-        sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
-    tmpfs_size : int, optional, default: None
-        The value for the size (in MiB) of the tmpfs mount for this step.
-        This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
-        memory allocated for this step.
-    tmpfs_path : str, optional, default /metaflow_temp
-        Path to tmpfs mount for this step.
-    persistent_volume_claims : Dict[str, str], optional, default None
-        A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
-        volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
-    shared_memory: int, optional
-        Shared memory size (in MiB) required for this step
-    port: int, optional
-        Port number to specify in the Kubernetes job object
-    compute_pool : str, optional, default None
-        Compute pool to be used for for this step.
-        If not specified, any accessible compute pool within the perimeter is used.
-    hostname_resolution_timeout: int, default 10 * 60
-        Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
-        Only applicable when @parallel is used.
-    """
-    def __init__(self, attributes = None, statically_defined = False):
-        ...
-    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
-        ...
-    def package_init(self, flow, step_name, environment):
-        ...
-    def runtime_init(self, flow, graph, package, run_id):
-        ...
-    def runtime_task_created(self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context):
-        ...
-    def runtime_step_cli(self, cli_args, retry_count, max_user_code_retries, ubf_context):
-        ...
-    def task_pre_step(self, step_name, task_datastore, metadata, run_id, task_id, flow, graph, retry_count, max_retries, ubf_context, inputs):
-        ...
-    def task_finished(self, step_name, flow, graph, is_task_ok, retry_count, max_retries):
-        ...
-    ...
-
-class Airflow(object, metaclass=type):
-    def __init__(self, name, graph, flow, code_package_sha, code_package_url, metadata, flow_datastore, environment, event_logger, monitor, production_token, tags = None, namespace = None, username = None, max_workers = None, worker_pool = None, description = None, file_path = None, workflow_timeout = None, is_paused_upon_creation = True):
-        ...
-    @classmethod
-    def get_existing_deployment(cls, name, flow_datastore):
-        ...
-    @classmethod
-    def get_token_path(cls, name):
-        ...
-    @classmethod
-    def save_deployment_token(cls, owner, name, token, flow_datastore):
-        ...
-    def compile(self):
-        ...
-    ...
-
-class AirflowException(metaflow.exception.MetaflowException, metaclass=type):
-    def __init__(self, msg):
-        ...
-    ...
-
-class NotSupportedException(metaflow.exception.MetaflowException, metaclass=type):
-    ...
+from ...metaflow_current import current as current
+from ..._vendor import click as click
+from ...exception import MetaflowException as MetaflowException
+from ...exception import MetaflowInternalError as MetaflowInternalError
+from ..aws.step_functions.production_token import load_token as load_token
+from ..aws.step_functions.production_token import new_token as new_token
+from ..aws.step_functions.production_token import store_token as store_token
+from ..kubernetes.kubernetes_decorator import KubernetesDecorator as KubernetesDecorator
+from .airflow import Airflow as Airflow
+from .exception import AirflowException as AirflowException
+from .exception import NotSupportedException as NotSupportedException
 
 class IncorrectProductionToken(metaflow.exception.MetaflowException, metaclass=type):
     ...

metaflow-stubs/plugins/airflow/airflow_decorator.pyi

@@ -1,36 +1,17 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.502360                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.621553                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
+import metaflow
 import typing
 if typing.TYPE_CHECKING:
     import metaflow.decorators
 
-class MetaDatum(tuple, metaclass=type):
-    """
-    MetaDatum(field, value, type, tags)
-    """
-    @staticmethod
-    def __new__(_cls, field, value, type, tags):
-        """
-        Create new instance of MetaDatum(field, value, type, tags)
-        """
-        ...
-    def __repr__(self):
-        """
-        Return a nicely formatted representation string
-        """
-        ...
-    def __getnewargs__(self):
-        """
-        Return self as a plain tuple.  Used by copy and pickle.
-        """
-        ...
-    ...
+from ...metadata_provider.metadata import MetaDatum as MetaDatum
 
 TASK_ID_XCOM_KEY: str
 

metaflow-stubs/plugins/airflow/airflow_utils.pyi

@@ -1,12 +1,13 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.501775                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.621209                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
 
+
 TASK_ID_XCOM_KEY: str
 
 FOREACH_CARDINALITY_XCOM_KEY: str
@@ -98,7 +99,7 @@ def get_metaflow_kubernetes_operator():
     ...
 
 class AirflowTask(object, metaclass=type):
-    def __init__(self, name, operator_type = "kubernetes", flow_name = None, is_mapper_node = False, flow_contains_foreach = False):
+    def __init__(self, name, operator_type = 'kubernetes', flow_name = None, is_mapper_node = False, flow_contains_foreach = False):
         ...
     @property
     def is_mapper_node(self):

metaflow-stubs/plugins/airflow/exception.pyi

@@ -1,21 +1,17 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.502622                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.621757                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
+import metaflow
 import typing
 if typing.TYPE_CHECKING:
     import metaflow.exception
 
-class MetaflowException(Exception, metaclass=type):
-    def __init__(self, msg = "", lineno = None):
-        ...
-    def __str__(self):
-        ...
-    ...
+from ...exception import MetaflowException as MetaflowException
 
 class AirflowException(metaflow.exception.MetaflowException, metaclass=type):
     def __init__(self, msg):

metaflow-stubs/plugins/airflow/sensors/__init__.pyi

@@ -1,104 +1,17 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.502956                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.622039                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
-import typing
-if typing.TYPE_CHECKING:
-    import metaflow.plugins.airflow.sensors.base_sensor
 
-class ExternalTaskSensorDecorator(metaflow.plugins.airflow.sensors.base_sensor.AirflowSensorDecorator, metaclass=type):
-    """
-    The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
-    This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
-    
-    Parameters
-    ----------
-    timeout : int
-        Time, in seconds before the task times out and fails. (Default: 3600)
-    poke_interval : int
-        Time in seconds that the job should wait in between each try. (Default: 60)
-    mode : str
-        How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
-    exponential_backoff : bool
-        allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
-    pool : str
-        the slot pool this task should run in,
-        slot pools are a way to limit concurrency for certain tasks. (Default:None)
-    soft_fail : bool
-        Set to true to mark the task as SKIPPED on failure. (Default: False)
-    name : str
-        Name of the sensor on Airflow
-    description : str
-        Description of sensor in the Airflow UI
-    external_dag_id : str
-        The dag_id that contains the task you want to wait for.
-    external_task_ids : List[str]
-        The list of task_ids that you want to wait for.
-        If None (default value) the sensor waits for the DAG. (Default: None)
-    allowed_states : List[str]
-        Iterable of allowed states, (Default: ['success'])
-    failed_states : List[str]
-        Iterable of failed or dis-allowed states. (Default: None)
-    execution_delta : datetime.timedelta
-        time difference with the previous execution to look at,
-        the default is the same logical date as the current task or DAG. (Default: None)
-    check_existence: bool
-        Set to True to check if the external task exists or check if
-        the DAG to wait for exists. (Default: True)
-    """
-    def serialize_operator_args(self):
-        ...
-    def validate(self, flow):
-        ...
-    ...
-
-class S3KeySensorDecorator(metaflow.plugins.airflow.sensors.base_sensor.AirflowSensorDecorator, metaclass=type):
-    """
-    The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
-    before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
-    and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
-    added as a flow decorators. Adding more than one decorator will ensure that `start` step
-    starts only after all sensors finish.
-    
-    Parameters
-    ----------
-    timeout : int
-        Time, in seconds before the task times out and fails. (Default: 3600)
-    poke_interval : int
-        Time in seconds that the job should wait in between each try. (Default: 60)
-    mode : str
-        How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
-    exponential_backoff : bool
-        allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
-    pool : str
-        the slot pool this task should run in,
-        slot pools are a way to limit concurrency for certain tasks. (Default:None)
-    soft_fail : bool
-        Set to true to mark the task as SKIPPED on failure. (Default: False)
-    name : str
-        Name of the sensor on Airflow
-    description : str
-        Description of sensor in the Airflow UI
-    bucket_key : Union[str, List[str]]
-        The key(s) being waited on. Supports full s3:// style url or relative path from root level.
-        When it's specified as a full s3:// url, please leave `bucket_name` as None
-    bucket_name : str
-        Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
-        When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
-    wildcard_match : bool
-        whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
-    aws_conn_id : str
-        a reference to the s3 connection on Airflow. (Default: None)
-    verify : bool
-        Whether or not to verify SSL certificates for S3 connection. (Default: None)
-    """
-    def validate(self, flow):
-        ...
-    ...
+from . import base_sensor as base_sensor
+from . import external_task_sensor as external_task_sensor
+from .external_task_sensor import ExternalTaskSensorDecorator as ExternalTaskSensorDecorator
+from . import s3_sensor as s3_sensor
+from .s3_sensor import S3KeySensorDecorator as S3KeySensorDecorator
 
 SUPPORTED_SENSORS: list
 

metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi

@@ -1,40 +1,19 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.558984                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.644264                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
+import metaflow
 import typing
 if typing.TYPE_CHECKING:
     import metaflow.decorators
-    import metaflow.exception
 
-class AirflowException(metaflow.exception.MetaflowException, metaclass=type):
-    def __init__(self, msg):
-        ...
-    ...
-
-class AirflowTask(object, metaclass=type):
-    def __init__(self, name, operator_type = "kubernetes", flow_name = None, is_mapper_node = False, flow_contains_foreach = False):
-        ...
-    @property
-    def is_mapper_node(self):
-        ...
-    def set_operator_args(self, **kwargs):
-        ...
-    def to_dict(self):
-        ...
-    @classmethod
-    def from_dict(cls, task_dict, flow_name = None, flow_contains_foreach = False):
-        ...
-    def to_task(self):
-        ...
-    ...
-
-def id_creator(val, hash_len):
-    ...
+from ..exception import AirflowException as AirflowException
+from ..airflow_utils import AirflowTask as AirflowTask
+from ..airflow_utils import id_creator as id_creator
 
 TASK_ID_HASH_LEN: int
 

metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi

@@ -1,50 +1,19 @@
-##################################################################################
-#                       Auto-generated Metaflow stub file                        #
-# MF version: 2.12.28                                                            #
-# Generated on 2024-11-01T10:21:04.559482                                        #
-##################################################################################
+######################################################################################################
+#                                 Auto-generated Metaflow stub file                                  #
+# MF version: 2.12.29                                                                                #
+# Generated on 2024-11-07T22:19:34.644568                                                            #
+######################################################################################################
 
 from __future__ import annotations
 
+import metaflow
 import typing
 if typing.TYPE_CHECKING:
-    import metaflow.decorators
     import metaflow.plugins.airflow.sensors.base_sensor
-    import metaflow.exception
 
-class AirflowSensorDecorator(metaflow.decorators.FlowDecorator, metaclass=type):
-    """
-    Base class for all Airflow sensor decorators.
-    """
-    def __init__(self, *args, **kwargs):
-        ...
-    def serialize_operator_args(self):
-        """
-        Subclasses will parse the decorator arguments to
-        Airflow task serializable arguments.
-        """
-        ...
-    def create_task(self):
-        ...
-    def validate(self, flow):
-        """
-        Validate if the arguments for the sensor are correct.
-        """
-        ...
-    def flow_init(self, flow, graph, environment, flow_datastore, metadata, logger, echo, options):
-        ...
-    ...
-
-class SensorNames(object, metaclass=type):
-    @classmethod
-    def get_supported_sensors(cls):
-        ...
-    ...
-
-class AirflowException(metaflow.exception.MetaflowException, metaclass=type):
-    def __init__(self, msg):
-        ...
-    ...
+from .base_sensor import AirflowSensorDecorator as AirflowSensorDecorator
+from ..airflow_utils import SensorNames as SensorNames
+from ..exception import AirflowException as AirflowException
 
 AIRFLOW_STATES: dict