support string conversion for scalars that contain a schema

Signed-off-by: redartera <[email protected]>
flyteorg · Jun 7, 2024 · e0b8b6f · e0b8b6f
1 parent 7521acb
commit e0b8b6f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/flytekit/interaction/string_literals.py b/flytekit/interaction/string_literals.py
@@ -37,6 +37,8 @@ def scalar_to_string(scalar: Scalar) -> typing.Any:
         return scalar.error.message
     if scalar.structured_dataset:
         return scalar.structured_dataset.uri
+    if scalar.schema:
+        return scalar.schema.uri
     if scalar.blob:
         return scalar.blob.uri
     if scalar.binary:

diff --git a/tests/flytekit/unit/interaction/test_string_literals.py b/tests/flytekit/unit/interaction/test_string_literals.py
@@ -22,10 +22,11 @@
     Primitive,
     Scalar,
     StructuredDataset,
+    Schema,
     Union,
     Void,
 )
-from flytekit.models.types import Error, LiteralType, SimpleType
+from flytekit.models.types import Error, LiteralType, SimpleType, SchemaType
 
 
 def test_primitive_to_string():
@@ -64,6 +65,9 @@ def test_scalar_to_string():
     scalar = Scalar(structured_dataset=StructuredDataset(uri="uri"))
     assert scalar_to_string(scalar) == "uri"
 
+    scalar = Scalar(schema=Schema(uri="schema_uri", type=SchemaType(columns=[])))
+    assert scalar_to_string(scalar) == "schema_uri"
+
     scalar = Scalar(
         blob=Blob(
             metadata=BlobMetadata(BlobType(format="", dimensionality=BlobType.BlobDimensionality.SINGLE)), uri="uri"