def test_bad_conversion():
orig = FlyteSchema[kwtypes(my_custom=bool)]
lt = TypeEngine.to_literal_type(orig)
# Make a not real column type
lt.schema.columns[0]._type = 15
with pytest.raises(ValueError):
TypeEngine.guess_python_type(lt)
def test_guessing_containers():
b = model_types.LiteralType(simple=model_types.SimpleType.BOOLEAN)
lt = model_types.LiteralType(collection_type=b)
pt = TypeEngine.guess_python_type(lt)
assert pt == typing.List[bool]
dur = model_types.LiteralType(simple=model_types.SimpleType.DURATION)
lt = model_types.LiteralType(map_value_type=dur)
pt = TypeEngine.guess_python_type(lt)
assert pt == typing.Dict[str, timedelta]
def test_bad_tag():
# Will not be able to load this
with pytest.raises(ValueError):
lt = LiteralType(simple=SimpleType.STRUCT, metadata={"pb_type": "bad.tag"})
TypeEngine.guess_python_type(lt)
# Doesn't match pb field key
with pytest.raises(ValueError):
lt = LiteralType(simple=SimpleType.STRUCT, metadata={})
TypeEngine.guess_python_type(lt)
def test_guessing_basic():
b = model_types.LiteralType(simple=model_types.SimpleType.BOOLEAN)
pt = TypeEngine.guess_python_type(b)
assert pt is bool
lt = model_types.LiteralType(simple=model_types.SimpleType.INTEGER)
pt = TypeEngine.guess_python_type(lt)
assert pt is int
lt = model_types.LiteralType(simple=model_types.SimpleType.STRING)
pt = TypeEngine.guess_python_type(lt)
assert pt is str
lt = model_types.LiteralType(simple=model_types.SimpleType.DURATION)
pt = TypeEngine.guess_python_type(lt)
assert pt is timedelta
lt = model_types.LiteralType(simple=model_types.SimpleType.DATETIME)
pt = TypeEngine.guess_python_type(lt)
assert pt is datetime.datetime
lt = model_types.LiteralType(simple=model_types.SimpleType.FLOAT)
pt = TypeEngine.guess_python_type(lt)
assert pt is float
lt = model_types.LiteralType(simple=model_types.SimpleType.NONE)
pt = TypeEngine.guess_python_type(lt)
assert pt is None
def test_jsondc_schemaize():
lt = TypeEngine.to_literal_type(Foo)
pt = TypeEngine.guess_python_type(lt)
# When postponed annotations are enabled, dataclass_json will not work and we'll end up with a
# schemaless generic.
# This test basically tests the broken behavior. Remove this test if
# https://github.com/lovasoa/marshmallow_dataclass/issues/13 is ever fixed.
assert pt is dict
def test_pb_guess_python_type():
artifact_tag = catalog_pb2.CatalogArtifactTag(artifact_id="artifact_1", name="artifact_name")
x = {"a": artifact_tag}
lt = TypeEngine.to_literal_type(catalog_pb2.CatalogArtifactTag)
gt = TypeEngine.guess_python_type(lt)
assert gt == catalog_pb2.CatalogArtifactTag
ctx = FlyteContextManager.current_context()
lm = TypeEngine.dict_to_literal_map(ctx, x, {"a": gt})
pv = TypeEngine.to_python_value(ctx, lm.literals["a"], gt)
assert pv == artifact_tag
def test_remaining_prims():
orig = FlyteSchema[kwtypes(my_dt=datetime, my_td=timedelta, my_b=bool)]
lt = TypeEngine.to_literal_type(orig)
pt = TypeEngine.guess_python_type(lt)
lt2 = TypeEngine.to_literal_type(pt)
assert lt == lt2
def test_schema_back_and_forth():
orig = FlyteSchema[kwtypes(TrackId=int, Name=str)]
lt = TypeEngine.to_literal_type(orig)
pt = TypeEngine.guess_python_type(lt)
lt2 = TypeEngine.to_literal_type(pt)
assert lt == lt2
