Adds type information to flytekit.models in repr (flyteorg#2389)

Signed-off-by: Thomas J. Fan <[email protected]>
jasonlai1218 · May 6, 2024 · 0869786 · 0869786
1 parent 8d258c4
commit 0869786
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 3 deletions.
diff --git a/flytekit/models/common.py b/flytekit/models/common.py
@@ -61,7 +61,8 @@ def short_string(self):
         :rtype: Text
         """
         literal_str = re.sub(r"\s+", " ", str(self.to_flyte_idl())).strip()
-        return f"<FlyteLiteral {literal_str}>"
+        type_str = type(self).__name__
+        return f"<FlyteLiteral({type_str}) {literal_str}>"
 
     def verbose_string(self):
         """

diff --git a/tests/flytekit/unit/core/test_promise.py b/tests/flytekit/unit/core/test_promise.py
@@ -84,7 +84,7 @@ def wf(i: int, j: int):
 
     # without providing the _inputs_not_allowed or _ignorable_inputs, all inputs to lp become required,
     # which is incorrect
-    with pytest.raises(FlyteAssertion, match="Missing input `i` type `<FlyteLiteral simple: INTEGER>`"):
+    with pytest.raises(FlyteAssertion, match=r"Missing input `i` type `<FlyteLiteral\(LiteralType\) simple: INTEGER>`"):
         create_and_link_node_from_remote(ctx, lp)
 
     # Even if j is not provided it will default

diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py
@@ -1701,7 +1701,7 @@ def wf2(a: typing.Union[int, str]) -> typing.Union[int, str]:
         match=re.escape(
             "Error encountered while executing 'wf2':\n"
             f"  Failed to convert inputs of task '{prefix}tests.flytekit.unit.core.test_type_hints.t2':\n"
-            '  Cannot convert from <FlyteLiteral scalar { union { value { scalar { primitive { string_value: "2" } } } '
+            '  Cannot convert from <FlyteLiteral(Literal) scalar { union { value { scalar { primitive { string_value: "2" } } } '
             'type { simple: STRING structure { tag: "str" } } } }> to typing.Union[float, dict] (using tag str)'
         ),
     ):

diff --git a/tests/flytekit/unit/models/test_common.py b/tests/flytekit/unit/models/test_common.py
@@ -103,3 +103,10 @@ def test_auth_role_empty():
     x = obj.to_flyte_idl()
     y = _common.AuthRole.from_flyte_idl(x)
     assert y == obj
+
+
+def test_short_string_raw_output_data_config():
+    """"""
+    obj = _common.RawOutputDataConfig("s3://bucket")
+    assert "FlyteLiteral(RawOutputDataConfig)" in obj.short_string()
+    assert "FlyteLiteral(RawOutputDataConfig)" in repr(obj)