From fc65e97bc60bf1719772f8a1708f32cab5bbb239 Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Tue, 10 Aug 2021 12:54:42 -0700 Subject: [PATCH] Protobuf transformer did not consider empty generic (#587) * Protobuf transformer did not consider empty generic - Zero value protobufs, become empty json (or empty dictionaries). - empty dictionaries currently result in an exception, when it should result in empty proto - This needs a simple change in the condition Signed-off-by: Ketan Umare * fix typo Signed-off-by: Ketan Umare --- flytekit/core/type_engine.py | 7 +++---- tests/flytekit/unit/core/test_type_engine.py | 7 +++++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index df175579c0..8cf7be6402 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -270,8 +270,8 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp return Literal(scalar=Scalar(generic=struct)) def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: - if not (lv and lv.scalar and lv.scalar.generic): - raise AssertionError("Can only covert a generic literal to a Protobuf") + if not (lv and lv.scalar and lv.scalar.generic is not None): + raise AssertionError("Can only convert a generic literal to a Protobuf") pb_obj = expected_python_type() dictionary = _MessageToDict(lv.scalar.generic) @@ -797,6 +797,7 @@ def _register_default_type_transformers(): TypeEngine.register(PathLikeTransformer()) TypeEngine.register(BinaryIOTransformer()) TypeEngine.register(EnumTransformer()) + TypeEngine.register(ProtobufTransformer()) # inner type is. Also unsupported are typing's Tuples. Even though you can look inside them, Flyte's type system # doesn't support these currently. @@ -810,5 +811,3 @@ def _register_default_type_transformers(): _register_default_type_transformers() - -TypeEngine.register(ProtobufTransformer()) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 1aecbe7151..a53a2f99cb 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -198,6 +198,13 @@ def test_protos(): with pytest.raises(AssertionError): TypeEngine.to_python_value(ctx, l0, errors_pb2.ContainerError) + default_proto = errors_pb2.ContainerError() + lit = TypeEngine.to_literal(ctx, default_proto, errors_pb2.ContainerError, lt) + assert lit.scalar + assert lit.scalar.generic is not None + new_python_val = TypeEngine.to_python_value(ctx, lit, errors_pb2.ContainerError) + assert new_python_val == default_proto + def test_guessing_basic(): b = model_types.LiteralType(simple=model_types.SimpleType.BOOLEAN)