flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

{union/_cli → flyte/cli}/_params.py

@@ -1,4 +1,3 @@
-import asyncio
 import dataclasses
 import datetime
 import enum
@@ -7,6 +6,7 @@ import importlib.util
 import json
 import os
 import pathlib
+import re
 import sys
 import typing
 import typing as t
@@ -14,24 +14,16 @@ from typing import get_args
 
 import rich_click as click
 import yaml
-from flyteidl.core.literals_pb2 import Literal
-from flyteidl.core.types_pb2 import BlobType, LiteralType, SimpleType
-
-from union._logging import logger
-from union.io import Dir, File
-from union.io.pickle.transformer import FlytePickleTransformer
-from union.storage._remote_fs import RemoteFSPathResolver
-from union.types import TypeEngine
-
-
-# ---------------------------------------------------
-# TODO replace these
-class ArtifactQuery:
-    pass
-
+from click import Context, Parameter
+from flyteidl2.core.interface_pb2 import Variable
+from flyteidl2.core.literals_pb2 import Literal
+from flyteidl2.core.types_pb2 import BlobType, LiteralType, SimpleType
+from google.protobuf.json_format import MessageToDict
+from mashumaro.codecs.json import JSONEncoder
 
-def is_remote(v: str) -> bool:
-    return False
+from flyte._logging import logger
+from flyte.io import Dir, File
+from flyte.types._pickle import FlytePickleTransformer
 
 
 class StructuredDataset:
@@ -43,26 +35,6 @@ class StructuredDataset:
 # ---------------------------------------------------
 
 
-def is_pydantic_basemodel(python_type: typing.Type) -> bool:
-    """
-    Checks if the python type is a pydantic BaseModel
-    """
-    try:
-        import pydantic  # noqa: F401
-    except ImportError:
-        return False
-    else:
-        try:
-            from pydantic import BaseModel as BaseModelV2
-            from pydantic.v1 import BaseModel as BaseModelV1
-
-            return issubclass(python_type, BaseModelV1) or issubclass(python_type, BaseModelV2)
-        except ImportError:
-            from pydantic import BaseModel
-
-        return issubclass(python_type, BaseModel)
-
-
 def key_value_callback(_: typing.Any, param: str, values: typing.List[str]) -> typing.Optional[typing.Dict[str, str]]:
     """
     Callback for click to parse key-value pairs.
@@ -97,20 +69,19 @@ def labels_callback(_: typing.Any, param: str, values: typing.List[str]) -> typi
 class DirParamType(click.ParamType):
     name = "directory path"
 
+    def get_metavar(self, param: Parameter, ctx: Context) -> str | None:
+        return "Remote Dir Path"
+
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
+        from flyte.storage import is_remote
 
-        # set remote_directory to false if running pyflyte run locally. This makes sure that the original
-        # directory is used and not a random one.
-        remote_directory = None if getattr(ctx.obj, "is_remote", False) else False
         if not is_remote(value):
             p = pathlib.Path(value)
             if not p.exists() or not p.is_dir():
                 raise click.BadParameter(f"parameter should be a valid flytedirectory path, {value}")
-        return Dir(path=value, remote_directory=remote_directory)
+        return Dir(path=value)
 
 
 class StructuredDatasetParamType(click.ParamType):
@@ -120,11 +91,12 @@ class StructuredDatasetParamType(click.ParamType):
 
     name = "structured dataset path (dir/file)"
 
+    def get_metavar(self, param: Parameter, ctx: Context) -> str | None:
+        return "Remote parquet URI"
+
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
         if isinstance(value, str):
             return StructuredDataset(uri=value)
         elif isinstance(value, StructuredDataset):
@@ -135,26 +107,27 @@ class StructuredDatasetParamType(click.ParamType):
 class FileParamType(click.ParamType):
     name = "file path"
 
+    def get_metavar(self, param: Parameter, ctx: Context) -> str | None:
+        return "Remote File Path"
+
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
-        # set remote_directory to false if running pyflyte run locally. This makes sure that the original
-        # file is used and not a random one.
-        remote_path = None if getattr(ctx.obj, "is_remote", False) else False
+        from flyte.storage import is_remote
+
         if not is_remote(value):
             p = pathlib.Path(value)
             if not p.exists() or not p.is_file():
                 raise click.BadParameter(f"parameter should be a valid file path, {value}")
-        return File(path=value, remote_path=remote_path)
+            raise click.BadParameter(f"Only remote paths are supported currently, {value}")
+        return File.from_existing_remote(value)
 
 
 class PickleParamType(click.ParamType):
     name = "pickle"
 
-    def get_metavar(self, param: click.Parameter) -> t.Optional[str]:
-        return "Python Object <Module>:<Object>"
+    def get_metavar(self, param: "Parameter", ctx) -> t.Optional[str]:
+        return "Python Object Instance <Module>:<Object>"
 
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
@@ -163,7 +136,7 @@ class PickleParamType(click.ParamType):
             return value
         parts = value.split(":")
         if len(parts) != 2:
-            if ctx and ctx.obj and ctx.obj.verbose > 0:
+            if ctx and ctx.obj and ctx.obj.log_level >= 10:  # DEBUG level
                 click.echo(f"Did not receive a string in the expected format <MODULE>:<VAR>, falling back to: {value}")
             return value
         try:
@@ -179,15 +152,15 @@ class PickleParamType(click.ParamType):
 class JSONIteratorParamType(click.ParamType):
     name = "json iterator"
 
+    def get_metavar(self, param: Parameter, ctx: Context) -> str | None:
+        return "JSON Value"
+
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
         return value
 
 
-import re
-
-
 def parse_iso8601_duration(iso_duration: str) -> datetime.timedelta:
     pattern = re.compile(
         r"^P"  # Starts with 'P'
@@ -211,10 +184,10 @@ def parse_human_durations(text: str) -> list[datetime.timedelta]:
     durations = []
 
     for part in raw_parts:
-        part = part.strip().lower()
+        new_part = part.strip().lower()
 
         # Match 1:24 or :45
-        m_colon = re.match(r"^(?:(\d+):)?(\d+)$", part)
+        m_colon = re.match(r"^(?:(\d+):)?(\d+)$", new_part)
         if m_colon:
             minutes = int(m_colon.group(1)) if m_colon.group(1) else 0
             seconds = int(m_colon.group(2))
@@ -222,7 +195,7 @@ def parse_human_durations(text: str) -> list[datetime.timedelta]:
             continue
 
         # Match "10 days", "1 minute", etc.
-        m_units = re.match(r"^(\d+)\s*(day|hour|minute|second)s?$", part)
+        m_units = re.match(r"^(\d+)\s*(day|hour|minute|second)s?$", new_part)
         if m_units:
             value = int(m_units.group(1))
             unit = m_units.group(2)
@@ -270,9 +243,6 @@ class DateTimeType(click.DateTime):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
-
         if isinstance(value, str) and " " in value:
             import re
 
@@ -300,11 +270,12 @@ class DateTimeType(click.DateTime):
 class DurationParamType(click.ParamType):
     name = "[1:24 | :22 | 1 minute | 10 days | ...]"
 
+    def get_metavar(self, param: Parameter, ctx: Context) -> str | None:
+        return "ISO8601 duration"
+
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Duration type.")
         return parse_duration(value)
@@ -318,8 +289,6 @@ class EnumParamType(click.Choice):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> enum.Enum:
-        if isinstance(value, ArtifactQuery):
-            return value
         if isinstance(value, self._enum_type):
             return value
         return self._enum_type(super().convert(value, param, ctx))
@@ -330,16 +299,20 @@ class UnionParamType(click.ParamType):
     A composite type that allows for multiple types to be specified. This is used for union types.
     """
 
-    def __init__(self, types: typing.List[click.ParamType]):
+    def __init__(self, types: typing.List[click.ParamType | None]):
         super().__init__()
         self._types = self._sort_precedence(types)
+        self.name = "|".join([t.name for t in self._types if t is not None])
+        self.optional = False
+        if None in types:
+            self.name = f"Optional[{self.name}]"
+            self.optional = True
 
-    @property
-    def name(self) -> str:
-        return "|".join([t.name for t in self._types])
+    def get_metavar(self, param: Parameter, ctx: typing.Optional[click.Context]) -> str | None:
+        return self.name
 
     @staticmethod
-    def _sort_precedence(tp: typing.List[click.ParamType]) -> typing.List[click.ParamType]:
+    def _sort_precedence(tp: typing.List[click.ParamType | None]) -> typing.List[click.ParamType]:
         unprocessed = []
         str_types = []
         others = []
@@ -350,7 +323,7 @@ class UnionParamType(click.ParamType):
                 str_types.append(p)
             else:
                 others.append(p)
-        return others + str_types + unprocessed
+        return others + str_types + unprocessed  # type: ignore
 
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
@@ -359,10 +332,10 @@ class UnionParamType(click.ParamType):
         Important to implement NoneType / Optional.
         Also could we just determine the click types from the python types
         """
-        if isinstance(value, ArtifactQuery):
-            return value
         for p in self._types:
             try:
+                if p is None and value is None:
+                    return None
                 return p.convert(value, param, ctx)
             except Exception as e:
                 logger.debug(f"Ignoring conversion error for type {p} trying other variants in Union. Error: {e}")
@@ -386,7 +359,7 @@ class JsonParamType(click.ParamType):
                 # We failed to load the json, so we'll try to load it as a file
                 if os.path.exists(value):
                     # if the value is a yaml file, we'll try to load it as yaml
-                    if value.endswith(".yaml") or value.endswith(".yml"):
+                    if value.endswith((".yaml", "yml")):
                         with open(value, "r") as f:
                             return yaml.safe_load(f)
                     with open(value, "r") as f:
@@ -398,8 +371,6 @@ class JsonParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-        if isinstance(value, ArtifactQuery):
-            return value
         if value is None:
             raise click.BadParameter("None value cannot be converted to a Json type.")
 
@@ -407,68 +378,38 @@ class JsonParamType(click.ParamType):
 
         # We compare the origin type because the json parsed value for list or dict is always a list or dict without
         # the covariant type information.
-        if type(parsed_value) == typing.get_origin(self._python_type) or type(parsed_value) == self._python_type:
+        if type(parsed_value) is typing.get_origin(self._python_type) or type(parsed_value) is self._python_type:
             # Indexing the return value of get_args will raise an error for native dict and list types.
             # We don't support native list/dict types with nested dataclasses.
             if get_args(self._python_type) == ():
                 return parsed_value
             elif isinstance(parsed_value, list) and dataclasses.is_dataclass(get_args(self._python_type)[0]):
                 j = JsonParamType(get_args(self._python_type)[0])
-                return [j.convert(v, param, ctx) for v in parsed_value]
+                # turn object back into json string
+                return [j.convert(json.dumps(v), param, ctx) for v in parsed_value]
             elif isinstance(parsed_value, dict) and dataclasses.is_dataclass(get_args(self._python_type)[1]):
                 j = JsonParamType(get_args(self._python_type)[1])
-                return {k: j.convert(v, param, ctx) for k, v in parsed_value.items()}
+                # turn object back into json string
+                return {k: j.convert(json.dumps(v), param, ctx) for k, v in parsed_value.items()}
 
             return parsed_value
 
-        if is_pydantic_basemodel(self._python_type):
-            """
-            This function supports backward compatibility for the Pydantic v1 plugin.
-            If the class is a Pydantic BaseModel, it attempts to parse JSON input using
-            the appropriate version of Pydantic (v1 or v2).
-            """
-            try:
-                if importlib.util.find_spec("pydantic.v1") is not None:
-                    from pydantic import BaseModel as BaseModelV2
-
-                    if issubclass(self._python_type, BaseModelV2):
-                        return self._python_type.model_validate_json(
-                            json.dumps(parsed_value), strict=False, context={"deserialize": True}
-                        )
-            except ImportError:
-                pass
-
-            # The behavior of the Pydantic v1 plugin.
-            return self._python_type.parse_raw(json.dumps(parsed_value))
-        return None
-
+        from pydantic import BaseModel
 
-def modify_literal_uris(lit: Literal):
-    """
-    Modifies the literal object recursively to replace the URIs with the native paths.
-    """
-    if lit.collection:
-        for l in lit.collection.literals:
-            modify_literal_uris(l)
-    elif lit.map:
-        for k, v in lit.map.literals.items():
-            modify_literal_uris(v)
-    elif lit.scalar:
-        if lit.scalar.blob and lit.scalar.blob.uri and lit.scalar.blob.uri.startswith(RemoteFSPathResolver.protocol):
-            lit.scalar.blob._uri = RemoteFSPathResolver.resolve_remote_path(lit.scalar.blob.uri)
-        elif lit.scalar.union:
-            modify_literal_uris(lit.scalar.union.value)
-        elif (
-            lit.scalar.structured_dataset
-            and lit.scalar.structured_dataset.uri
-            and lit.scalar.structured_dataset.uri.startswith(RemoteFSPathResolver.protocol)
-        ):
-            lit.scalar.structured_dataset._uri = RemoteFSPathResolver.resolve_remote_path(
-                lit.scalar.structured_dataset.uri
+        if issubclass(self._python_type, BaseModel):
+            return typing.cast(BaseModel, self._python_type).model_validate_json(
+                json.dumps(parsed_value), strict=False, context={"deserialize": True}
             )
+        elif dataclasses.is_dataclass(self._python_type):
+            from mashumaro.codecs.json import JSONDecoder
 
+            decoder = JSONDecoder(self._python_type)
+            return decoder.decode(value)
 
-SIMPLE_TYPE_CONVERTER: typing.Dict[SimpleType, click.ParamType] = {
+        return parsed_value
+
+
+SIMPLE_TYPE_CONVERTER = {
     SimpleType.FLOAT: click.FLOAT,
     SimpleType.INTEGER: click.INT,
     SimpleType.STRING: click.STRING,
@@ -482,47 +423,53 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
     """
     Converts a Flyte LiteralType given a python_type to a click.ParamType
     """
-    if lt.simple:
+    if lt.HasField("simple"):
         if lt.simple == SimpleType.STRUCT:
             ct = JsonParamType(python_type)
             ct.name = f"JSON object {python_type.__name__}"
             return ct
         if lt.simple in SIMPLE_TYPE_CONVERTER:
             return SIMPLE_TYPE_CONVERTER[lt.simple]
-        raise NotImplementedError(f"Type {lt.simple} is not supported in pyflyte run")
-
-    if lt.enum_type:
-        return EnumParamType(python_type)  # type: ignore
+        raise NotImplementedError(f"Type {lt.simple} is not supported in `flyte run`")
 
-    if lt.structured_dataset_type:
+    if lt.HasField("structured_dataset_type"):
         return StructuredDatasetParamType()
 
-    if lt.collection_type or lt.map_value_type:
+    if lt.HasField("collection_type") or lt.HasField("map_value_type"):
         ct = JsonParamType(python_type)
-        if lt.collection_type:
+        if lt.HasField("collection_type"):
             ct.name = "json list"
         else:
             ct.name = "json dictionary"
         return ct
 
-    if lt.blob:
+    if lt.HasField("blob"):
         if lt.blob.dimensionality == BlobType.BlobDimensionality.SINGLE:
             if lt.blob.format == FlytePickleTransformer.PYTHON_PICKLE_FORMAT:
                 return PickleParamType()
+            # TODO: Add JSONIteratorTransformer
             # elif lt.blob.format == JSONIteratorTransformer.JSON_ITERATOR_FORMAT:
             #     return JSONIteratorParamType()
             return FileParamType()
         return DirParamType()
 
-    if lt.union_type:
+    if lt.HasField("union_type"):
+        python_args = typing.get_args(python_type)
+        if len(python_args) == 0:
+            return PickleParamType()
         cts = []
         for i in range(len(lt.union_type.variants)):
             variant = lt.union_type.variants[i]
-            variant_python_type = typing.get_args(python_type)[i]
-            ct = literal_type_to_click_type(variant, variant_python_type)
-            cts.append(ct)
+            variant_python_type = python_args[i]
+            if variant_python_type is type(None):
+                cts.append(None)
+            else:
+                cts.append(literal_type_to_click_type(variant, variant_python_type))
         return UnionParamType(cts)
 
+    if lt.HasField("enum_type"):
+        return EnumParamType(python_type)  # type: ignore
+
     return click.UNPROCESSED
 
 
@@ -533,9 +480,7 @@ class FlyteLiteralConverter(object):
         self,
         literal_type: LiteralType,
         python_type: typing.Type,
-        is_remote: bool,
     ):
-        self._is_remote = is_remote
         self._literal_type = literal_type
         self._python_type = python_type
         self._click_type = literal_type_to_click_type(literal_type, python_type)
@@ -547,29 +492,24 @@ class FlyteLiteralConverter(object):
     def is_bool(self) -> bool:
         return self.click_type == click.BOOL
 
+    def is_optional(self) -> bool:
+        return isinstance(self.click_type, UnionParamType) and self.click_type.optional
+
     def convert(
         self, ctx: click.Context, param: typing.Optional[click.Parameter], value: typing.Any
     ) -> typing.Union[Literal, typing.Any]:
         """
-        Convert the value to a Flyte Literal or a python native type. This is used by click to convert the input.
+        Convert the value to a python native type. This is used by click to convert the input.
         """
-        if isinstance(value, ArtifactQuery):
-            return value
         try:
             # If the expected Python type is datetime.date, adjust the value to date
             if self._python_type is datetime.date:
+                if value is None:
+                    return None
                 # Click produces datetime, so converting to date to avoid type mismatch error
                 value = value.date()
-            # If the input matches the default value in the launch plan, serialization can be skipped.
-            if param and value == param.default:
-                return None
 
-            # If this is used for remote execution, then we need to convert it back to a python native type
-            if not self._is_remote:
-                return value
-
-            lit = asyncio.run(TypeEngine.to_literal(value, self._python_type, self._literal_type))
-            return lit
+            return value
         except click.BadParameter:
             raise
         except Exception as e:
@@ -577,3 +517,62 @@ class FlyteLiteralConverter(object):
                 f"Failed to convert param: {param if param else 'NA'}, value: {value} to type: {self._python_type}."
                 f" Reason {e}"
             ) from e
+
+
+def to_click_option(
+    input_name: str,
+    literal_var: Variable,
+    python_type: typing.Type,
+    default_val: typing.Any,
+) -> click.Option:
+    """
+    This handles converting workflow input types to supported click parameters with callbacks to initialize
+    the input values to their expected types.
+    """
+    from flyteidl2.core.types_pb2 import SimpleType
+
+    if input_name != input_name.lower():
+        # Click does not support uppercase option names: https://github.com/pallets/click/issues/837
+        raise ValueError(f"Workflow input name must be lowercase: {input_name!r}")
+
+    literal_converter = FlyteLiteralConverter(
+        literal_type=literal_var.type,
+        python_type=python_type,
+    )
+
+    if literal_converter.is_bool() and not default_val:
+        default_val = False
+
+    description_extra = ""
+    if literal_var.type.simple == SimpleType.STRUCT:
+        if default_val:
+            # pydantic v2
+            if hasattr(default_val, "model_dump_json"):
+                default_val = default_val.model_dump_json()
+            else:
+                encoder = JSONEncoder(python_type)
+                default_val = encoder.encode(default_val)
+        if literal_var.type.metadata:
+            description_extra = f": {MessageToDict(literal_var.type.metadata)}"
+
+    required = False if default_val is not None else True
+    is_flag: typing.Optional[bool] = None
+    param_decls = [f"--{input_name}"]
+    if literal_converter.is_bool():
+        required = False
+        is_flag = True
+        if default_val is True:
+            param_decls = [f"--{input_name}/--no-{input_name}"]
+    if literal_converter.is_optional():
+        required = False
+
+    return click.Option(
+        param_decls=param_decls,
+        type=literal_converter.click_type,
+        is_flag=is_flag,
+        default=default_val,
+        show_default=True,
+        required=required,
+        help=literal_var.description + description_extra,
+        callback=literal_converter.convert,
+    )