Fix 0.0 value floats (flyteorg#452)

Signed-off-by: wild-endeavor <[email protected]> Signed-off-by: Max Hoffman <[email protected]>
dolthub · Apr 29, 2021 · 86af4cf · 86af4cf
1 parent fa063ff
commit 86af4cf
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
@@ -530,9 +530,9 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
 
 
 def _check_and_covert_float(lv: Literal) -> float:
-    if lv.scalar.primitive.float_value:
+    if lv.scalar.primitive.float_value is not None:
         return lv.scalar.primitive.float_value
-    elif lv.scalar.primitive.integer:
+    elif lv.scalar.primitive.integer is not None:
         return float(lv.scalar.primitive.integer)
     raise RuntimeError(f"Cannot convert literal {lv} to float")
 

diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py
@@ -197,3 +197,13 @@ def test_protos():
     l0 = Literal(scalar=Scalar(primitive=Primitive(integer=4)))
     with pytest.raises(AssertionError):
         TypeEngine.to_python_value(ctx, l0, errors_pb2.ContainerError)
+
+
+def test_zero_floats():
+    ctx = FlyteContext.current_context()
+
+    l0 = Literal(scalar=Scalar(primitive=Primitive(integer=0)))
+    l1 = Literal(scalar=Scalar(primitive=Primitive(float_value=0.0)))
+
+    assert TypeEngine.to_python_value(ctx, l0, float) == 0
+    assert TypeEngine.to_python_value(ctx, l1, float) == 0