Skip to main content

Creating a Custom Type Transformer

When a Python class is not recognized by Flyte, TypeEngine eventually selects FlytePickleTransformer; to give the class a stable Flyte representation, implement TypeTransformer and register it before Flyte builds the task interface.

import asyncio

from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
from flyteidl.core import literals_pb2, types_pb2


class Color:
def __init__(self, name: str):
self.name = name


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

def get_literal_type(self, t):
return types_pb2.LiteralType(simple=types_pb2.SimpleType.STRING)

async def to_literal(self, python_val, python_type, expected):
if not isinstance(python_val, Color):
raise TypeTransformerFailedError(
f"Expected Color object, received {type(python_val)}"
)

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

async def to_python_value(self, lv, expected_python_type):
if not lv.HasField("scalar") or not lv.scalar.HasField("primitive"):
raise TypeTransformerFailedError(f"Expected primitive literal, received {lv}")
if not lv.scalar.primitive.HasField("string_value"):
raise TypeTransformerFailedError(f"Expected string literal, received {lv}")

return Color(lv.scalar.primitive.string_value)

def guess_python_type(self, literal_type):
if (
literal_type.HasField("simple")
and literal_type.simple == types_pb2.SimpleType.STRING
):
return Color
raise ValueError(f"Cannot guess Color from {literal_type}")


TypeEngine.register(ColorTransformer())

literal_type = TypeEngine.to_literal_type(Color)
literal = asyncio.run(
TypeEngine.to_literal(Color("blue"), Color, literal_type)
)
restored = asyncio.run(TypeEngine.to_python_value(literal, Color))
assert restored.name == "blue"

Implement the transformer contract

flyte.types re-exports TypeEngine, TypeTransformer, and TypeTransformerFailedError from the type-engine module. A transformer is parameterized by the Python type it handles and is initialized with a display name and that type:

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

The base constructor stores these values as name, python_type, and type_assertions_enabled. Type assertions are enabled by default. Before calling a transformer's to_literal, TypeEngine.to_literal calls transformer.assert_type(python_type, python_val) when that flag is enabled. The base implementation rejects a value that is not an instance of the declared type with TypeTransformerFailedError. A custom transformer can override assert_type for more specific validation, or pass enable_type_assertions=False to super().__init__ when the standard check is unsuitable.

Three methods are abstract and required:

  • get_literal_type(self, t) maps the Python annotation to a FlyteIDL LiteralType.
  • async to_literal(self, python_val, python_type, expected) serializes a value to a FlyteIDL Literal.
  • async to_python_value(self, lv, expected_python_type) deserializes a Literal to the requested Python type.

The to_literal method receives both the runtime value and the declared python_type. The TypeTransformer contract specifically tells implementers to use the passed declared type rather than relying only on type(python_val). The expected argument is the literal type generated for the declaration; blob transformers such as FileTransformer use it as part of the conversion contract while constructing their literal.

TypeTransformerFailedError inherits from TypeError, AssertionError, and ValueError. Raise it when the value or literal has the wrong shape, as FileTransformer does when it receives a non-blob literal or a multipart blob where a single blob is expected.

Choose a Flyte literal representation

The representation in get_literal_type must agree with the representation emitted by to_literal and accepted by to_python_value. The ColorTransformer declares a SimpleType.STRING and emits a primitive string. For a URI-backed value, the built-in FileTransformer instead declares a single-part BlobType:

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

def get_literal_type(self, t: Type[File]) -> types_pb2.LiteralType:
return types_pb2.LiteralType(
blob=types_pb2.BlobType(
format="",
dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE,
)
)

Its to_literal places the path and format in the blob literal and preserves an optional hash:

return literals_pb2.Literal(
scalar=literals_pb2.Scalar(
blob=literals_pb2.Blob(
metadata=literals_pb2.BlobMetadata(
type=types_pb2.BlobType(
format=python_val.format,
dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE,
)
),
uri=python_val.path,
)
),
hash=python_val.hash if python_val.hash else None,
)

FileTransformer does not perform I/O; its docstring assigns I/O responsibility to the user. A custom blob-like transformer should therefore keep its literal conversion separate from any upload or download behavior unless it explicitly needs that behavior.

For a binary representation, TypeTransformer.from_binary_idl provides a reusable MessagePack decoding path when the literal contains a Binary tagged with MESSAGEPACK. The built-in dataclass transformer and dictionary transformer use this pattern and cache MessagePackDecoder instances. A custom transformer can override from_binary_idl when it needs a different binary format or validation.

The base to_html implementation returns str(python_val). Override it when the custom type needs a different HTML representation. TypeEngine.to_html also recognizes a Renderable object attached to an Annotated type and delegates to its to_html method.

Register the transformer

Register one transformer instance after its class definition:

TypeEngine.register(ColorTransformer())

TypeEngine.register adds the transformer's python_type to the class-level registry. It accepts additional_types when the same transformer should handle multiple concrete Python types:

TypeEngine.register(ColorTransformer(), additional_types=[SpecialColor])

Every primary type and alias must be absent from the registry. A duplicate raises ValueError and reports both the existing and attempted transformer names. This makes registration timing important: import the module that performs registration before constructing interfaces or converting values. The built-in FileTransformer, DirTransformer, and pickle transformer all register at module import time.

For an already-created transformer, register_additional_type adds a mapping for another type:

transformer = ColorTransformer()
TypeEngine.register(transformer)
TypeEngine.register_additional_type(transformer, SpecialColor)

Without override=True, an existing additional-type mapping is left unchanged. Passing override=True intentionally replaces that mapping:

TypeEngine.register_additional_type(transformer, SpecialColor, override=True)

Use this method for a shared transformer that supports several implementations. The dataframe integration follows this approach by registering its shared DataFrameTransformerEngine for concrete dataframe handler types rather than creating an unrelated top-level transformer for each handler.

Exercise the engine directly

Use to_literal_type to create the interface type, then pass that type to the asynchronous conversion method. The reverse operation takes the literal and the expected Python type:

import asyncio

from flyte.types import TypeEngine


literal_type = TypeEngine.to_literal_type(Color)

literal = asyncio.run(
TypeEngine.to_literal(
python_val=Color("green"),
python_type=Color,
expected=literal_type,
)
)

value = asyncio.run(TypeEngine.to_python_value(literal, Color))
assert isinstance(value, Color)
assert value.name == "green"

TypeEngine.to_literal selects the transformer, performs its enabled type assertion, awaits to_literal, and then calls modify_literal_uris on the resulting literal. TypeEngine.to_python_value first unwraps an offloaded literal when offloaded_metadata is present, selects the transformer for expected_python_type, and awaits to_python_value.

The runtime uses the same calls at task boundaries. For outputs, _internal/runtime/convert.py calls to_literal_type and to_literal for every declared output. A TypeTransformerFailedError is converted into RuntimeDataValidationError with the output name and task name:

for (output_name, python_type), v in zip(interface.outputs.items(), o):
try:
lit = await TypeEngine.to_literal(
v, python_type, TypeEngine.to_literal_type(python_type)
)
named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
except TypeTransformerFailedError as e:
raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)

For inputs, the same runtime module creates a LiteralMap and calls TypeEngine.literal_map_to_kwargs. That method converts each literal using the native input type mapping. A transformer therefore needs both directions to work in a task: output serialization alone is not sufficient for a downstream task to receive the custom value.

Support reverse type inference

Implement guess_python_type when a LiteralType may need to be converted back into a Python annotation. The default TypeTransformer.guess_python_type raises ValueError, and TypeEngine.guess_python_type tries registered transformers until one accepts the literal. This path is used by interface helpers and by literal_map_to_kwargs when only Flyte variable types, rather than Python type hints, are available.

FileTransformer checks the blob shape and excludes the pickle format before returning File:

def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[File]:
if (
literal_type.HasField("blob")
and literal_type.blob.dimensionality
== types_pb2.BlobType.BlobDimensionality.SINGLE
and literal_type.blob.format != "PythonPickle"
):
return File
raise ValueError(f"Cannot guess python type from {literal_type}")

The ColorTransformer example uses the simpler SimpleType.STRING check. If multiple registered transformers accept the same literal shape, reverse inference can select the first matching transformer, so a custom literal type should carry enough distinguishing information for guess_python_type to reject unrelated literals.

Understand lookup and fallback behavior

TypeEngine.get_transformer checks registrations for the requested type and then applies several lookup rules:

  1. Annotated types are inspected for a TypeTransformer instance in their metadata; otherwise lookup continues with the underlying type.
  2. Enum subclasses use the built-in enum transformer.
  3. Generic aliases are checked as a complete registered type and then by their origin, such as list for a list alias.
  4. PEP 604 union types use the registered union transformer.
  5. Exact registrations and the method-resolution order are searched.
  6. Dataclasses use the built-in dataclass transformer only after the MRO search.
  7. An otherwise unknown type falls back to FlytePickleTransformer.

The late dataclass step allows a user transformer registered for a dataclass-like type to win during MRO lookup. Register custom handling before the type is used to build an interface.

Pickle fallback can serialize arbitrary Python objects, but the package documentation identifies the format as non-human-readable, not representable in the UI, and potentially inefficient for large datasets. A registered transformer such as ColorTransformer instead declares a normal Flyte literal and defines how to validate and reconstruct it.

Troubleshoot custom transformers

The transformer is never selected

Check that registration ran before TypeEngine.to_literal_type or task interface construction. Inspect the registry through TypeEngine.get_available_transformers() and confirm that the exact Python class, or the intended additional type, is present. Also check that a generic annotation is registered at the type or origin that lookup examines.

Registration raises ValueError

The primary type or one of the additional_types is already registered. Do not call register repeatedly for the same type. If replacing an alias is intentional, use register_additional_type(..., override=True); TypeEngine.register itself does not override an existing mapping.

Deserialization rejects the literal

Validate the literal shape before reading nested protobuf fields. FileTransformer.to_python_value checks for a blob and verifies single-part dimensionality before constructing File; a custom transformer should similarly reject a scalar, collection, blob, or binary representation that it does not understand by raising TypeTransformerFailedError.

Reverse conversion fails

Implement guess_python_type if code receives only a LiteralType. Without an override, the base method raises ValueError, and TypeEngine.guess_python_type may ultimately report that no transformer could reverse the type.

A value reaches pickle unexpectedly

The engine did not find a matching registration. Verify the import that performs registration, the declared annotation used by the task, and any alias or generic-origin mapping. The runtime uses the declared python_type, not merely the runtime class of the value.

Round-trip checklist

  • Register the transformer before interface construction.
  • Implement all three abstract methods with their exact asynchronous signatures.
  • Make get_literal_type, to_literal, and to_python_value agree on the protobuf shape.
  • Decide whether the default assert_type is sufficient; override it or disable assertions only when needed.
  • Reject malformed literals with TypeTransformerFailedError.
  • Test TypeEngine.to_literal_type, to_literal, and to_python_value together.
  • Add a guess_python_type test if interfaces or literal maps may be reconstructed from Flyte types.
  • Test declared annotations separately from runtime values, including any Annotated, optional, or generic forms the transformer is intended to support.
  • Test the task-boundary path, because runtime input conversion uses literal_map_to_kwargs and output conversion wraps transformer failures as runtime data-validation errors.