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 <[email protected]> * fix typo Signed-off-by: Ketan Umare <[email protected]>
flyteorg · Aug 10, 2021 · fc65e97 · fc65e97
1 parent e8449c7
commit fc65e97
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git 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
@@ -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)