Unions, Enums, and Optional Types

In this SDK, branching logic and optionality are handled through standard Python type hints. The UnionTransformer and EnumTransformer classes in flyte.types._type_engine manage the conversion between Python's flexible type system and Flyte's structured IDL.

Optional Types

The most common use of branching logic is the Optional type. In Python, Optional[T] is a shorthand for Union[T, None]. The UnionTransformer specifically detects this pattern using the is_optional_type static method:

# src/flyte/types/_type_engine.py

@staticmethod
def is_optional_type(t: Type[Any]) -> bool:
    return _is_union_type(t) and type(None) in get_args(t)

When a task input is marked as Optional, Flyte allows the value to be missing or explicitly set to None. This is frequently used for parameters with default values.

Example: Optional Task Inputs

In examples/basics/union_type_simple.py, a task handles an optional string:

@env.task
def process_union_input(input_val: Union[str, None]) -> str:
    if input_val is None:
        return "Input was None!"
    return f"Input was a string: {input_val}"

Union Types

UnionTransformer allows a single input or output to accept multiple distinct types, such as Union[int, str] or Union[MyDataclass, str].

Structural Matching and Ambiguity

When converting a Python value to a Flyte literal, the UnionTransformer iterates through all possible subtypes defined in the Union. It attempts to use the TypeEngine to find a transformer for each subtype.

A critical implementation detail is how the SDK handles ambiguity. If a value could validly represent more than one type in the Union (for example, two different Dataclasses that have the same field names and types), the UnionTransformer will raise a TypeError.

# src/flyte/types/_type_engine.py (UnionTransformer.to_literal)

if is_ambiguous:
    raise TypeError(
        f"Ambiguous choice of variant for union type.\n"
        f"Potential types: {potential_types}\n"
        "These types are structurally the same, because it's attributes have the same names and associated types."
    )

To avoid this, ensure that the types within a Union are structurally distinct or use unique field names in your Dataclasses and Pydantic models.

Enums

The EnumTransformer enables the use of enum.Enum for categorical inputs. This provides a way to restrict task inputs to a predefined set of valid strings.

String-Valued Enums

Flyte primarily supports string-valued enums. The EnumTransformer enforces this during the conversion process:

# src/flyte/types/_type_engine.py (EnumTransformer.get_literal_type)

values = [v.value for v in t]  # type: ignore
if not isinstance(values[0], str):
    raise TypeTransformerFailedError("Only EnumTypes with name of value are supported")

Names vs. Values

By default, the SDK serializes enums using their member names, not their values. For example, if you have Color.RED = "red-value", Flyte will store the string "RED" in the underlying literal.

# src/flyte/types/_type_engine.py (EnumTransformer.to_literal)

return Literal(scalar=Scalar(primitive=Primitive(string_value=python_val.name)))

Example: Defining Enums

As seen in examples/basics/enums.py, you can define enums using standard Python syntax:

class Color(str, enum.Enum):
    RED = "red-value"
    GREEN = "green-value"
    BLUE = "blue-value"

@env.task
async def standalone_enum_echo(color: Color) -> str:
    return f"color={color.name}({color.value})"

Literal Types

The SDK also supports typing.Literal for cases where you want to restrict a value to a specific set of constants without defining a full Enum class. The EnumTransformer handles these by treating them as a special LITERAL_ENUM.

In examples/basics/all_types.py, typing.Literal is used for simple categorical flags:

Intensity = Literal["low", "medium", "high"]

@env.task
async def process_custom_types(my_literal: Intensity) -> Dict[str, str]:
    return {"literal": f"Intensity: {my_literal}"}

When using typing.Literal, the EnumTransformer uses the literal values themselves rather than member names, as there are no member names associated with a typing.Literal.

Integration with Complex Types

Unions and Enums are recursively supported within more complex structures like Pydantic models and Dataclasses. When the PydanticTransformer or DataclassTransformer encounters a field typed as a Union or Enum, it delegates the serialization and deserialization to the UnionTransformer and EnumTransformer respectively.

This allows for deeply nested structures with branching logic:

class ShirtOrder(BaseModel):
    color: Color  # Enum
    size: Size    # Enum
    quantity: int

class FullOrder(BaseModel):
    items: list[ShirtOrder]
    notes: Optional[str] = None  # Optional/Union

In this example, the TypeEngine coordinates multiple transformers to ensure every field is correctly mapped to the Flyte IDL.