Skip to content

fix: EnumTransformer.to_literal accepts a string matching an enum value#3439

Open
1fanwang wants to merge 3 commits into
flyteorg:masterfrom
1fanwang:enum-to-literal-accepts-string
Open

fix: EnumTransformer.to_literal accepts a string matching an enum value#3439
1fanwang wants to merge 3 commits into
flyteorg:masterfrom
1fanwang:enum-to-literal-accepts-string

Conversation

@1fanwang

@1fanwang 1fanwang commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Why are the changes needed?

EnumTransformer.assert_type accepts a raw string that matches one of the enum's values (for instance an enum default supplied as a string), but EnumTransformer.to_literal rejects that same string with Expected an enum. So a matching-string value passes type-checking and then fails at serialization.

What changes were proposed in this pull request?

to_literal now accepts a string matching one of the enum's values and serializes it to the corresponding literal, consistent with assert_type. A non-matching string is still rejected.

How was this patch tested?

Added test_enum_to_literal_accepts_matching_string — it fails on the current code with TypeTransformerFailedError: Expected an enum and passes with the fix.

Check all the applicable boxes

  • I updated the documentation accordingly.
  • All new and existing tests passed.
  • All commits are signed-off.

@1fanwang 1fanwang force-pushed the enum-to-literal-accepts-string branch from f995050 to 11c9eaf Compare June 19, 2026 18:56
Comment thread flytekit/core/type_engine.py Outdated
Comment thread flytekit/core/type_engine.py Outdated
res = None
res_type = None
# Result of the ``str`` variant, if it matched, used to disambiguate a bare-string value.
str_match: typing.Optional[typing.Tuple[Literal, LiteralType]] = None

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding str_match can make the logic here more complex. Could we update this function to:

    async def async_to_literal(
        self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType
    ) -> typing.Union[Literal, asyncio.Future]:
        python_type = get_underlying_type(python_type)

        matches: list[tuple[type, Literal, LiteralType]] = []
        for i, t in enumerate(get_args(python_type)):
            try:
                trans: TypeTransformer[T] = TypeEngine.get_transformer(t)
                if isinstance(trans, AsyncTypeTransformer):
                    res = await trans.async_to_literal(ctx, python_val, t, expected.union_type.variants[i])
                else:
                    res = trans.to_literal(ctx, python_val, t, expected.union_type.variants[i])
                res_type = _add_tag_to_type(trans.get_literal_type(t), trans.name)
                matches.append((t, res, res_type))
            except Exception as e:
                logger.debug(f"UnionTransformer failed attempt to convert from {python_val} to {t} error: {e}")

        if not matches:
            raise TypeTransformerFailedError(f"Cannot convert from {python_val} to {python_type}")

        if len(matches) > 1:
            # NOTE: More than one variant accepted the value. Pick the one whose type is
            # exactly the same as python_val's type. For example, given Union[str, Color]
            # and python_val="red", both str and Color accept it (Color accepts a matching
            # string), but str is the exact match so it wins.
            exact = [m for m in matches if type(python_val) is m[0]]
            if len(exact) != 1:
                raise TypeError(
                    f"Ambiguous choice of variant for union type.\n"
                    f"Potential types: {[m[0] for m in matches]}\n"
                    "These types are structurally the same, because it's attributes have the same names and associated types."
                )
            matches = exact

        _, res, res_type = matches[0]
        return Literal(scalar=Scalar(union=Union(value=res, stored_type=res_type)))

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 39c76dc. Heads-up: exact-type-wins now resolves Union[int, UnsignedInt] + 3 to int instead of raising, so I updated test_union_custom_transformer_sanity_check to assert that (it matches test_union_custom_transformer, which already expects int to win for 3).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correction to the above: the general exact-type tie-break broke test_ambiguous_union_transformer_to_literal and the sanity_check (in 2 files) -- it silently resolves structurally-identical variants (Union[A, B], Union[int, UnsignedInt]) that must stay ambiguous, since they are indistinguishable on read-back. Narrowed it to just the str-vs-enum case in 033d1cc, which keeps that guard. Happy to drop the guard and update those tests instead if you prefer.

Comment thread tests/flytekit/unit/core/test_type_engine.py Outdated
Comment thread tests/flytekit/unit/core/test_type_engine.py Outdated
Comment thread tests/flytekit/unit/core/test_type_engine.py Outdated
1fanwang added 2 commits June 30, 2026 00:07
- Simplify EnumTransformer.to_literal: unify the enum/string paths through
  a single value extraction
- UnionTransformer: replace the str-specific disambiguation with a general
  "exact type wins" tie-break when multiple variants match
- Tests: de-parametrize the matching-string case, use docstrings, and
  update the custom-transformer sanity check for the exact-type tie-break

Signed-off-by: 1fanwang <1fannnw@gmail.com>
…ty guard

The general exact-type tie-break removed the guard that makes structurally-
identical union variants (e.g. Union[A, B] dataclasses, Union[int, UnsignedInt])
raise on to_literal -- they would then be indistinguishable on read-back. Keep
that guard; only a bare string prefers the str variant over an enum-by-value.

Signed-off-by: 1fanwang <1fannnw@gmail.com>

@machichima machichima left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thank you. Just a small nit left

Comment on lines +1084 to +1085
# Accept a raw string matching one of the enum's values: assert_type already allows it (e.g. an enum
# default supplied as a string), so to_literal must too, else such a value type-checks but fails to serialize.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Those details should be on PR description rather than comment. Let's simply to:

Suggested change
# Accept a raw string matching one of the enum's values: assert_type already allows it (e.g. an enum
# default supplied as a string), so to_literal must too, else such a value type-checks but fails to serialize.
# Accept either an enum member or a raw string matching one of its values.

@machichima

Copy link
Copy Markdown
Member

Hi @1fanwang ,
I merged your CI fixed #3448, could you please rebase/merge master?
Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants