def test_enum_type():
t = TypeEngine.to_literal_type(Color)
assert t is not None
assert t.enum_type is not None
assert t.enum_type.values
assert t.enum_type.values == [c.value for c in Color]
ctx = FlyteContextManager.current_context()
lv = TypeEngine.to_literal(ctx, Color.RED, Color, TypeEngine.to_literal_type(Color))
assert lv
assert lv.scalar
assert lv.scalar.primitive.string_value == "red"
v = TypeEngine.to_python_value(ctx, lv, Color)
assert v
assert v == Color.RED
v = TypeEngine.to_python_value(ctx, lv, str)
assert v
assert v == "red"
with pytest.raises(ValueError):
TypeEngine.to_python_value(ctx, Literal(scalar=Scalar(primitive=Primitive(string_value=str(Color.RED)))), Color)
with pytest.raises(ValueError):
TypeEngine.to_python_value(ctx, Literal(scalar=Scalar(primitive=Primitive(string_value="bad"))), Color)
with pytest.raises(AssertionError):
TypeEngine.to_literal_type(UnsupportedEnumValues)
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
def test_list_transformer():
l0 = Literal(scalar=Scalar(primitive=Primitive(integer=3)))
l1 = Literal(scalar=Scalar(primitive=Primitive(integer=4)))
lc = LiteralCollection(literals=[l0, l1])
lit = Literal(collection=lc)
ctx = FlyteContext.current_context()
xx = TypeEngine.to_python_value(ctx, lit, typing.List[int])
assert xx == [3, 4]
def convert(
self,
value: typing.Any,
python_type_hint: typing.Optional[typing.Type] = None) -> Scalar:
if self.primitive_type:
return Scalar(primitive=Primitive(**{self.primitive_type: value}))
if self.scalar_type:
return Scalar(**{self.scalar_type: value})
raise NotImplementedError("Not implemented yet!")
def test_protos():
ctx = FlyteContext.current_context()
pb = errors_pb2.ContainerError(code="code", message="message")
lt = TypeEngine.to_literal_type(errors_pb2.ContainerError)
assert lt.simple == SimpleType.STRUCT
assert lt.metadata["pb_type"] == "flyteidl.core.errors_pb2.ContainerError"
lit = TypeEngine.to_literal(ctx, pb, errors_pb2.ContainerError, lt)
new_python_val = TypeEngine.to_python_value(ctx, lit, errors_pb2.ContainerError)
assert new_python_val == pb
# Test error
l0 = Literal(scalar=Scalar(primitive=Primitive(integer=4)))
with pytest.raises(AssertionError):
TypeEngine.to_python_value(ctx, l0, errors_pb2.ContainerError)
def to_literal(
self,
ctx: FlyteContext,
python_val: Union[FlyteFile, FlyteSchema, str],
python_type: Type[GreatExpectationsType],
expected: LiteralType,
) -> Literal:
datatype = GreatExpectationsTypeTransformer.get_config(python_type)[0]
if issubclass(datatype, FlyteSchema):
return FlyteSchemaTransformer().to_literal(ctx, python_val,
datatype, expected)
elif issubclass(datatype, FlyteFile):
return FlyteFilePathTransformer().to_literal(
ctx, python_val, datatype, expected)
elif issubclass(datatype, str):
return Literal(scalar=Scalar(primitive=Primitive(
string_value=python_val)))
else:
raise TypeError(f"{datatype} is not a supported type")
def test_dict_transformer():
d = DictTransformer()
def assert_struct(lit: LiteralType):
assert lit is not None
assert lit.simple == SimpleType.STRUCT
def recursive_assert(lit: LiteralType,
expected: LiteralType,
expected_depth: int = 1,
curr_depth: int = 0):
assert curr_depth <= expected_depth
assert lit is not None
if lit.map_value_type is None:
assert lit == expected
return
recursive_assert(lit.map_value_type, expected, expected_depth,
curr_depth + 1)
# Type inference
assert_struct(d.get_literal_type(dict))
assert_struct(d.get_literal_type(typing.Dict[int, int]))
recursive_assert(d.get_literal_type(typing.Dict[str, str]),
LiteralType(simple=SimpleType.STRING))
recursive_assert(d.get_literal_type(typing.Dict[str, int]),
LiteralType(simple=SimpleType.INTEGER))
recursive_assert(d.get_literal_type(typing.Dict[str, datetime.datetime]),
LiteralType(simple=SimpleType.DATETIME))
recursive_assert(d.get_literal_type(typing.Dict[str, datetime.timedelta]),
LiteralType(simple=SimpleType.DURATION))
recursive_assert(d.get_literal_type(typing.Dict[str, dict]),
LiteralType(simple=SimpleType.STRUCT))
recursive_assert(
d.get_literal_type(typing.Dict[str, typing.Dict[str, str]]),
LiteralType(simple=SimpleType.STRING),
expected_depth=2,
)
recursive_assert(
d.get_literal_type(typing.Dict[str, typing.Dict[int, str]]),
LiteralType(simple=SimpleType.STRUCT),
expected_depth=2,
)
recursive_assert(
d.get_literal_type(typing.Dict[str, typing.Dict[str,
typing.Dict[str,
str]]]),
LiteralType(simple=SimpleType.STRING),
expected_depth=3,
)
recursive_assert(
d.get_literal_type(typing.Dict[str, typing.Dict[str,
typing.Dict[str,
dict]]]),
LiteralType(simple=SimpleType.STRUCT),
expected_depth=3,
)
recursive_assert(
d.get_literal_type(typing.Dict[str, typing.Dict[str,
typing.Dict[int,
dict]]]),
LiteralType(simple=SimpleType.STRUCT),
expected_depth=2,
)
ctx = FlyteContext.current_context()
lit = d.to_literal(ctx, {}, typing.Dict, LiteralType(SimpleType.STRUCT))
pv = d.to_python_value(ctx, lit, typing.Dict)
assert pv == {}
# Literal to python
with pytest.raises(TypeError):
d.to_python_value(
ctx, Literal(scalar=Scalar(primitive=Primitive(integer=10))), dict)
with pytest.raises(TypeError):
d.to_python_value(ctx, Literal(), dict)
with pytest.raises(TypeError):
d.to_python_value(ctx, Literal(map=LiteralMap(literals={"x": None})),
dict)
with pytest.raises(TypeError):
d.to_python_value(ctx, Literal(map=LiteralMap(literals={"x": None})),
typing.Dict[int, str])
d.to_python_value(
ctx,
Literal(map=LiteralMap(
literals={
"x": Literal(scalar=Scalar(primitive=Primitive(integer=1)))
})),
typing.Dict[str, int],
)
def _register_default_type_transformers():
TypeEngine.register(
SimpleTransformer(
"int",
int,
_primitives.Integer.to_flyte_literal_type(),
lambda x: Literal(scalar=Scalar(primitive=Primitive(integer=x))),
lambda x: x.scalar.primitive.integer,
))
TypeEngine.register(
SimpleTransformer(
"float",
float,
_primitives.Float.to_flyte_literal_type(),
lambda x: Literal(scalar=Scalar(primitive=Primitive(float_value=x))
),
_check_and_covert_float,
))
TypeEngine.register(
SimpleTransformer(
"bool",
bool,
_primitives.Boolean.to_flyte_literal_type(),
lambda x: Literal(scalar=Scalar(primitive=Primitive(boolean=x))),
lambda x: x.scalar.primitive.boolean,
))
TypeEngine.register(
SimpleTransformer(
"str",
str,
_primitives.String.to_flyte_literal_type(),
lambda x: Literal(scalar=Scalar(primitive=Primitive(string_value=x)
)),
lambda x: x.scalar.primitive.string_value,
))
TypeEngine.register(
SimpleTransformer(
"datetime",
_datetime.datetime,
_primitives.Datetime.to_flyte_literal_type(),
lambda x: Literal(scalar=Scalar(primitive=Primitive(datetime=x))),
lambda x: x.scalar.primitive.datetime,
))
TypeEngine.register(
SimpleTransformer(
"timedelta",
_datetime.timedelta,
_primitives.Timedelta.to_flyte_literal_type(),
lambda x: Literal(scalar=Scalar(primitive=Primitive(duration=x))),
lambda x: x.scalar.primitive.duration,
))
TypeEngine.register(
SimpleTransformer(
"none",
None,
_type_models.LiteralType(simple=_type_models.SimpleType.NONE),
lambda x: None,
lambda x: None,
))
TypeEngine.register(ListTransformer())
TypeEngine.register(DictTransformer())
TypeEngine.register(TextIOTransformer())
TypeEngine.register(PathLikeTransformer())
TypeEngine.register(BinaryIOTransformer())
# 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.
# Confusing note: typing.NamedTuple is in here even though task functions themselves can return them. We just mean
# that the return signature of a task can be a NamedTuple that contains another NamedTuple inside it.
# Also, it's not entirely true that Flyte IDL doesn't support tuples. We can always fake them as structs, but we'll
# hold off on doing that for now, as we may amend the IDL formally to support tuples.
TypeEngine.register(RestrictedType("non typed tuple", tuple))
TypeEngine.register(RestrictedType("non typed tuple", typing.Tuple))
TypeEngine.register(RestrictedType("named tuple", typing.NamedTuple))
def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal:
return Literal(scalar=Scalar(primitive=Primitive(string_value=python_val.value)))
def test_interface():
ctx = FlyteContextManager.current_context()
lt = TypeEngine.to_literal_type(pd.DataFrame)
df = pd.DataFrame({"name": ["Tom", "Joseph"], "age": [20, 22]})
annotated_sd_type = Annotated[StructuredDataset, kwtypes(name=str, age=int)]
df_literal_type = TypeEngine.to_literal_type(annotated_sd_type)
assert df_literal_type.structured_dataset_type is not None
assert len(df_literal_type.structured_dataset_type.columns) == 2
assert df_literal_type.structured_dataset_type.columns[0].name == "name"
assert df_literal_type.structured_dataset_type.columns[0].literal_type.simple is not None
assert df_literal_type.structured_dataset_type.columns[1].name == "age"
assert df_literal_type.structured_dataset_type.columns[1].literal_type.simple is not None
sd = annotated_sd_type(df)
sd_literal = TypeEngine.to_literal(ctx, sd, python_type=annotated_sd_type, expected=lt)
lm = {
"my_map": Literal(
map=LiteralMap(
literals={
"k1": Literal(scalar=Scalar(primitive=Primitive(string_value="v1"))),
"k2": Literal(scalar=Scalar(primitive=Primitive(string_value="2"))),
},
)
),
"my_list": Literal(
collection=LiteralCollection(
literals=[
Literal(scalar=Scalar(primitive=Primitive(integer=1))),
Literal(scalar=Scalar(primitive=Primitive(integer=2))),
Literal(scalar=Scalar(primitive=Primitive(integer=3))),
]
)
),
"val_a": Literal(scalar=Scalar(primitive=Primitive(integer=21828))),
"my_df": sd_literal,
}
variable_map = {
"my_map": interface_models.Variable(type=TypeEngine.to_literal_type(typing.Dict[str, str]), description=""),
"my_list": interface_models.Variable(type=TypeEngine.to_literal_type(typing.List[int]), description=""),
"val_a": interface_models.Variable(type=TypeEngine.to_literal_type(int), description=""),
"my_df": interface_models.Variable(type=df_literal_type, description=""),
}
lr = LiteralsResolver(lm, variable_map=variable_map, ctx=ctx)
assert lr._ctx is ctx
with pytest.raises(ValueError):
lr["not"] # noqa
with pytest.raises(ValueError):
lr.get_literal("not")
# Test that just using [] works, guessing from the Flyte type is invoked
result = lr["my_list"]
assert result == [1, 2, 3]
# Test that using get works, guessing from the Flyte type is invoked
result = lr.get("my_map")
assert result == {
"k1": "v1",
"k2": "2",
}
# Getting the literal will return the Literal object itself
assert lr.get_literal("my_df") is sd_literal
guessed_df = lr["my_df"]
# Based on guessing, so no column information
assert len(guessed_df.metadata.structured_dataset_type.columns) == 0
guessed_df_2 = lr["my_df"]
assert guessed_df is guessed_df_2
# Update type hints with the annotated type
lr.update_type_hints({"my_df": annotated_sd_type})
del lr._native_values["my_df"]
guessed_df = lr.get("my_df")
# Using the user specified type, so number of columns is correct.
assert len(guessed_df.metadata.structured_dataset_type.columns) == 2
from flytekit import kwtypes
from flytekit.core.context_manager import FlyteContextManager
from flytekit.core.type_engine import LiteralsResolver, TypeEngine
from flytekit.models import interface as interface_models
from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar
from flytekit.types.structured.structured_dataset import StructuredDataset
@pytest.mark.parametrize(
"literal_value,python_type,expected_python_value",
[
(
Literal(
collection=LiteralCollection(
literals=[
Literal(scalar=Scalar(primitive=Primitive(integer=1))),
Literal(scalar=Scalar(primitive=Primitive(integer=2))),
Literal(scalar=Scalar(primitive=Primitive(integer=3))),
]
)
),
typing.List[int],
[1, 2, 3],
),
(
Literal(
map=LiteralMap(
literals={
"k1": Literal(scalar=Scalar(primitive=Primitive(string_value="v1"))),
"k2": Literal(scalar=Scalar(primitive=Primitive(string_value="2"))),
},