diff --git a/metaflow/metaflow_config.py b/metaflow/metaflow_config.py index b13de89c986..75343a4fe08 100644 --- a/metaflow/metaflow_config.py +++ b/metaflow/metaflow_config.py @@ -1,6 +1,10 @@ import os import sys import types +import uuid +import datetime + +from typing import Dict, List, Union, Tuple as TTuple from metaflow.exception import MetaflowException from metaflow.metaflow_config_funcs import from_conf, get_validate_choice_fn @@ -615,6 +619,56 @@ def get_pinned_conda_libs(python_version, datastore_type): return pins +### +# Runner API type mappings +# Extensions can add custom Click parameter types via get_click_to_python_types +### +def get_click_to_python_types(): + """ + Returns the mapping from Click parameter types to Python types for Runner API. + Extensions can override this function to add custom type mappings. + """ + from metaflow._vendor.click.types import ( + BoolParamType, + Choice, + DateTime, + File, + FloatParamType, + IntParamType, + Path, + StringParamType, + Tuple, + UUIDParameterType, + ) + from metaflow.parameters import JSONTypeClass + from metaflow.includefile import FilePathClass + from metaflow.user_configs.config_options import ( + LocalFileInput, + MultipleTuple, + ConfigValue, + ) + + # Define JSON type for type hints + JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None] + + return { + StringParamType: str, + IntParamType: int, + FloatParamType: float, + BoolParamType: bool, + UUIDParameterType: uuid.UUID, + Path: str, + DateTime: datetime.datetime, + Tuple: tuple, + Choice: str, + File: str, + JSONTypeClass: JSON, + FilePathClass: str, + LocalFileInput: str, + MultipleTuple: TTuple[str, Union[JSON, ConfigValue]], + } + + # Check if there are extensions to Metaflow to load and override everything try: from metaflow.extension_support import get_modules @@ -650,6 +704,16 @@ def _new_get_pinned_conda_libs( if any(" " in x for x in o): raise ValueError("Decospecs cannot contain spaces") _TOGGLE_DECOSPECS.extend(o) + elif n == "get_click_to_python_types": + # Extension provides additional Click type mappings for Runner API + # Merge extension's types with base types + def _new_get_click_to_python_types(f1=globals()[n], f2=o): + d1 = f1() + d2 = f2() + d1.update(d2) + return d1 + + globals()[n] = _new_get_click_to_python_types elif not n.startswith("__") and not isinstance(o, types.ModuleType): globals()[n] = o # If DEFAULT_DECOSPECS is set, use that, else extrapolate from extensions diff --git a/metaflow/runner/click_api.py b/metaflow/runner/click_api.py index 202eff6dc05..cf53a550090 100644 --- a/metaflow/runner/click_api.py +++ b/metaflow/runner/click_api.py @@ -57,26 +57,14 @@ ) from metaflow.user_decorators.user_flow_decorator import FlowMutator +# Import Click type mappings from config (allows extensions to add custom types) +from metaflow.metaflow_config import get_click_to_python_types + +click_to_python_types = get_click_to_python_types() + # Define a recursive type alias for JSON JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None] -click_to_python_types = { - StringParamType: str, - IntParamType: int, - FloatParamType: float, - BoolParamType: bool, - UUIDParameterType: uuid.UUID, - Path: str, - DateTime: datetime.datetime, - Tuple: tuple, - Choice: str, - File: str, - JSONTypeClass: JSON, - FilePathClass: str, - LocalFileInput: str, - MultipleTuple: TTuple[str, Union[JSON, ConfigValue]], -} - def _method_sanity_check( possible_arg_params: TOrderedDict[str, click.Argument],