Creating Custom Type Transformers

Flyte's type system is designed to be extensible, allowing you to use domain-specific Python types in your tasks and workflows. By implementing a TypeTransformer, you define how a custom Python object is converted into a Flyte-understandable format (a Literal) and back again.

In this tutorial, you will build a PositiveInt type that validates integers are greater than zero and create a transformer to integrate it with Flyte.

Prerequisites

To follow this tutorial, you need the following imports from the Flyte SDK and its dependencies:

import typing
from typing import Type, Optional
from flyteidl2.core import literals_pb2, types_pb2
from flyte.types._type_engine import TypeTransformer, TypeEngine, TypeTransformerFailedError

Step 1: Define the Custom Python Type

First, create the Python class that represents your domain-specific data. In this example, PositiveInt wraps a standard integer but adds validation logic in its constructor.

# Defined in examples/type_transformers/my_transformer/src/my_transformer/custom_type.py

class PositiveInt:
    """A wrapper type that only accepts positive integers."""

    def __init__(self, value: int):
        if not isinstance(value, int):
            raise TypeError(f"Expected int, got {type(value).__name__}")
        if value <= 0:
            raise ValueError(f"Expected positive integer, got {value}")
        self._value = value

    @property
    def value(self) -> int:
        return self._value

    def __repr__(self) -> str:
        return f"PositiveInt({self._value})"

This class ensures that any instance of PositiveInt always contains a valid positive integer.

Step 2: Implement the Type Transformer

To make Flyte aware of PositiveInt, you must implement the TypeTransformer interface. This class handles the mapping between your Python type and Flyte's internal IDL (Interface Definition Language) types.

Define the Transformer Class

Inherit from TypeTransformer[PositiveInt] and initialize it with a name and the target Python type.

# Based on examples/type_transformers/my_transformer/src/my_transformer/transformer.py

class PositiveIntTransformer(TypeTransformer[PositiveInt]):
    def __init__(self):
        super().__init__(name="PositiveInt", t=PositiveInt)

Map to a Flyte Literal Type

Implement get_literal_type to tell Flyte which underlying primitive or complex type should represent your data. Since PositiveInt is essentially an integer, we use SimpleType.INTEGER.

    def get_literal_type(self, t: Type[PositiveInt]) -> types_pb2.LiteralType:
        return types_pb2.LiteralType(
            simple=types_pb2.SimpleType.INTEGER,
            structure=types_pb2.TypeStructure(tag="PositiveInt"),
        )

Convert Python to Flyte Literal

The to_literal method is called when a task returns a PositiveInt or receives one as an input. It converts the Python object into a literals_pb2.Literal.

    async def to_literal(
        self,
        python_val: PositiveInt,
        python_type: Type[PositiveInt],
        expected: types_pb2.LiteralType,
    ) -> literals_pb2.Literal:
        if not isinstance(python_val, PositiveInt):
            raise TypeTransformerFailedError(f"Expected PositiveInt, got {type(python_val).__name__}")

        return literals_pb2.Literal(
            scalar=literals_pb2.Scalar(primitive=literals_pb2.Primitive(integer=python_val.value))
        )

Convert Flyte Literal to Python

The to_python_value method performs the reverse operation, reconstructing your Python object from a Flyte Literal. This is where you can re-apply your custom validation.

    async def to_python_value(self, lv: literals_pb2.Literal, expected_python_type: Type[PositiveInt]) -> PositiveInt:
        if not lv.scalar or not lv.scalar.primitive:
            raise TypeTransformerFailedError("Missing scalar primitive")

        value = lv.scalar.primitive.integer
        try:
            return PositiveInt(value)
        except (TypeError, ValueError) as e:
            raise TypeTransformerFailedError(f"Validation failed: {e}")

Step 3: Register the Transformer

For Flyte to use your transformer, you must register it with the TypeEngine. This is typically done at the module level or in an __init__.py file to ensure it is registered before any tasks are defined.

TypeEngine.register(PositiveIntTransformer())

Note that TypeEngine.register will raise a ValueError if a transformer for the same type is already registered.

Step 4: Use the Custom Type in a Task

Once registered, you can use PositiveInt in your task signatures just like any standard Python type.

# Based on examples/type_transformers/main.py
from flyte.core.task import task

@task
def process_positive_int(x: PositiveInt) -> PositiveInt:
    print(f"Processing: {x}")
    return PositiveInt(x.value + 10)

When this task runs, Flyte uses your PositiveIntTransformer to serialize the input and deserialize the output.

Shortcut: Using SimpleTransformer

If your custom type is a simple wrapper around a primitive that doesn't require complex async logic, you can use SimpleTransformer to reduce boilerplate. It accepts lambda functions for the transformations.

# Example of using SimpleTransformer from flyte.types._type_engine
from flyte.types._type_engine import SimpleTransformer

def to_lit(val: PositiveInt) -> literals_pb2.Literal:
    return literals_pb2.Literal(scalar=literals_pb2.Scalar(primitive=literals_pb2.Primitive(integer=val.value)))

def from_lit(lv: literals_pb2.Literal) -> PositiveInt:
    return PositiveInt(lv.scalar.primitive.integer)

simple_transformer = SimpleTransformer(
    name="SimplePositiveInt",
    t=PositiveInt,
    lt=types_pb2.LiteralType(simple=types_pb2.SimpleType.INTEGER),
    to_literal_transformer=to_lit,
    from_literal_transformer=from_lit
)

TypeEngine.register(simple_transformer)

SimpleTransformer is ideal for types that map directly to Flyte's SimpleType primitives like INTEGER, FLOAT, or STRING.