def test_enum_to_execution():
@op(
config_schema={
"my_enum":
Field(Enum.from_python_enum(TestEnum),
is_required=False,
default_value="ONE")
})
def my_op(context):
return context.op_config["my_enum"]
@graph
def my_graph():
my_op()
my_job = my_graph.to_job()
result = my_job.execute_in_process()
assert result.success
assert result.output_for_node("my_op") == TestEnum.ONE
result = my_graph.execute_in_process(
run_config={"ops": {
"my_op": {
"config": {
"my_enum": "TWO"
}
}
}})
assert result.success
assert result.output_for_node("my_op") == TestEnum.TWO
def test_receive_processed_config_values():
class TestEnum(enum.Enum):
FOO = 1
BAR = 2
enum_conf_schema = {
"foo":
Field(Enum.from_python_enum(TestEnum),
is_required=False,
default_value="BAR")
}
@config_mapping(config_schema=enum_conf_schema)
def processed_config_mapping(outer_config):
return {"ops": {"my_op": {"config": {"foo": outer_config["foo"]}}}}
processed_result = my_graph.to_job(
config=processed_config_mapping).execute_in_process()
assert processed_result.success
assert processed_result.output_for_node("my_op") == TestEnum.BAR
@config_mapping(config_schema=enum_conf_schema,
receive_processed_config_values=False)
def unprocessed_config_mapping(outer_config):
return {"ops": {"my_op": {"config": {"foo": outer_config["foo"]}}}}
unprocessed_result = my_graph.to_job(
config=unprocessed_config_mapping).execute_in_process()
assert unprocessed_result.success
assert unprocessed_result.output_for_node("my_op") == "BAR"
def test_native_enum_dagster_enum_from_classmethod_default_value():
dagster_enum = Enum.from_python_enum(NativeEnum)
called = {}
@solid(config=Field(dagster_enum, is_required=False, default_value='BAR'))
def enum_name_as_str(context):
assert context.solid_config == NativeEnum.BAR
called['str'] = True
@solid(config=Field(dagster_enum,
is_required=False,
default_value=NativeEnum.BAR))
def enum_direct(context):
assert context.solid_config == NativeEnum.BAR
called['enum'] = True
@pipeline
def test():
enum_direct()
enum_name_as_str()
result = execute_pipeline(test)
assert result.success
assert called['str']
assert called['enum']
def test_native_enum_not_allowed_as_default_value():
dagster_enum = Enum.from_python_enum(NativeEnum)
with pytest.raises(DagsterInvalidDefinitionError) as exc_info:
@solid(config_schema=Field(dagster_enum, is_required=False, default_value=NativeEnum.BAR))
def _enum_direct(_):
pass
assert str(exc_info.value) == (
"You have passed into a python enum value as the default value into "
"of a config enum type NativeEnum. You must pass in the underlying string "
"represention as the default value. One of ['FOO', 'BAR']."
)
def test_list_enum_with_bad_default_value():
dagster_enum = Enum.from_python_enum(NativeEnum)
with pytest.raises(DagsterInvalidConfigError) as exc_info:
@solid(config=Field([dagster_enum], is_required=False, default_value=[NativeEnum.BAR]))
def _bad_enum_list(_):
pass
# This error message is bad
# https://github.com/dagster-io/dagster/issues/2339
assert 'Invalid default_value for Field.' in str(exc_info.value)
assert 'Error 1: Value at path root[0] for enum type NativeEnum must be a string' in str(
exc_info.value
)
def test_native_enum_dagster_enum_from_classmethod():
dagster_enum = Enum.from_python_enum(NativeEnum)
called = {}
@solid(config=dagster_enum)
def dagster_enum_me(context):
assert context.solid_config == NativeEnum.BAR
called['yup'] = True
pipeline_def = PipelineDefinition(
name='native_enum_dagster_pipeline', solid_defs=[dagster_enum_me]
)
result = execute_pipeline(pipeline_def, {'solids': {'dagster_enum_me': {'config': 'BAR'}}})
assert result.success
assert called['yup']
def test_native_enum_dagster_enum_from_classmethod():
dagster_enum = Enum.from_python_enum(NativeEnum)
called = {}
@solid(config_schema=dagster_enum)
def dagster_enum_me(context):
assert context.solid_config == NativeEnum.BAR
called["yup"] = True
pipeline_def = PipelineDefinition(
name="native_enum_dagster_pipeline", solid_defs=[dagster_enum_me]
)
result = execute_pipeline(pipeline_def, {"solids": {"dagster_enum_me": {"config": "BAR"}}})
assert result.success
assert called["yup"]
def test_dict_enum_with_default():
dagster_enum = Enum.from_python_enum(NativeEnum)
called = {}
@solid(config_schema={"enum": Field(dagster_enum, is_required=False, default_value="BAR")})
def enum_dict(context):
assert context.solid_config["enum"] == NativeEnum.BAR
called["yup"] = True
@pipeline
def enum_dict_pipeline():
enum_dict()
result = execute_pipeline(enum_dict_pipeline)
assert result.success
assert called["yup"]
def test_list_enum_with_default_value():
dagster_enum = Enum.from_python_enum(NativeEnum)
called = {}
@solid(config_schema=Field([dagster_enum], is_required=False, default_value=["BAR"]))
def enum_list(context):
assert context.solid_config == [NativeEnum.BAR]
called["yup"] = True
@pipeline
def enum_list_pipeline():
enum_list()
result = execute_pipeline(enum_list_pipeline)
assert result.success
assert called["yup"]
def test_native_enum_classmethod_creates_all_values():
dagster_enum = Enum.from_python_enum(NativeEnum)
for enum_value in NativeEnum:
assert enum_value is dagster_enum.to_python_value(enum_value.name)
def test_enum_config_mapping():
@op(
config_schema={
"my_enum":
Field(Enum.from_python_enum(TestEnum),
is_required=False,
default_value="ONE")
})
def my_op(context):
return context.op_config["my_enum"]
@graph
def my_graph():
my_op()
def _use_defaults_mapping(_):
return {}
use_defaults = my_graph.to_job(config=ConfigMapping(
config_fn=_use_defaults_mapping))
result = use_defaults.execute_in_process()
assert result.success
assert result.output_for_node("my_op") == TestEnum.ONE
def _override_defaults_mapping(_):
return {"ops": {"my_op": {"config": {"my_enum": "TWO"}}}}
override_defaults = my_graph.to_job(config=ConfigMapping(
config_fn=_override_defaults_mapping))
result = override_defaults.execute_in_process()
assert result.success
assert result.output_for_node("my_op") == TestEnum.TWO
def _ingest_config_mapping(x):
return {"ops": {"my_op": {"config": {"my_enum": x["my_field"]}}}}
default_config_mapping = ConfigMapping(
config_fn=_ingest_config_mapping,
config_schema=Shape({
"my_field":
Field(Enum.from_python_enum(TestEnum),
is_required=False,
default_value="TWO")
}),
receive_processed_config_values=False,
)
ingest_mapping = my_graph.to_job(config=default_config_mapping)
result = ingest_mapping.execute_in_process()
assert result.success
assert result.output_for_node("my_op") == TestEnum.TWO
no_default_config_mapping = ConfigMapping(
config_fn=_ingest_config_mapping,
config_schema=Shape({
"my_field":
Field(Enum.from_python_enum(TestEnum), is_required=True)
}),
receive_processed_config_values=False,
)
ingest_mapping_no_default = my_graph.to_job(
config=no_default_config_mapping)
result = ingest_mapping_no_default.execute_in_process(
run_config={"my_field": "TWO"})
assert result.success
assert result.output_for_node("my_op") == TestEnum.TWO
def _ingest_post_processed_config(x):
assert x["my_field"] == TestEnum.TWO
return {"ops": {"my_op": {"config": {"my_enum": "TWO"}}}}
config_mapping_with_preprocessing = ConfigMapping(
config_fn=_ingest_post_processed_config,
config_schema=Shape({
"my_field":
Field(Enum.from_python_enum(TestEnum), is_required=True)
}),
)
ingest_preprocessing = my_graph.to_job(
config=config_mapping_with_preprocessing)
result = ingest_preprocessing.execute_in_process(
run_config={"my_field": "TWO"})
assert result.success
assert result.output_for_node("my_op") == TestEnum.TWO
