Enums as dict keys are not supported by Config system

[danielgafni]

Dagster version

1.5.9

What's the issue?

I would like to have a config field which is a dictionary with limited options for keys. Specifically, I want the keys to be valid AWS instance types.

I'd expect the Enum type to work for this, but it doesn't.

What did you expect to happen?

The docs here mention Enums are supported in config fields, but they only work for dict values, not keys.

How to reproduce?

from enum import Enum
from typing import Dict
from dagster import materialize, asset, Config
from pydantic import Field

def test_enum_key():
    class MyEnum(Enum):
        FOO = "foo"
        BAR = "bar"

    class MyConfig(Config):
        baz: Dict[MyEnum, str] = Field(default_factory=lambda: {MyEnum.FOO: "foo"})

    @asset
    def my_asset(config: MyConfig):
        return config.baz

    materialize(
        [my_asset]
    )

test_enums_config.py:15: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.venv/lib/python3.11/site-packages/dagster/_core/decorator_utils.py:195: in wrapped_with_pre_call_fn
    return fn(*args, **kwargs)
.venv/lib/python3.11/site-packages/dagster/_core/decorator_utils.py:195: in wrapped_with_pre_call_fn
    return fn(*args, **kwargs)
.venv/lib/python3.11/site-packages/dagster/_core/decorator_utils.py:195: in wrapped_with_pre_call_fn
    return fn(*args, **kwargs)
.venv/lib/python3.11/site-packages/dagster/_core/decorator_utils.py:195: in wrapped_with_pre_call_fn
    return fn(*args, **kwargs)
.venv/lib/python3.11/site-packages/dagster/_core/decorator_utils.py:195: in wrapped_with_pre_call_fn
    return fn(*args, **kwargs)
.venv/lib/python3.11/site-packages/dagster/_core/definitions/decorators/asset_decorator.py:247: in asset
    return create_asset()(compute_fn)
.venv/lib/python3.11/site-packages/dagster/_core/definitions/decorators/asset_decorator.py:416: in __call__
    op = _Op(
.venv/lib/python3.11/site-packages/dagster/_core/definitions/decorators/op_decorator.py:108: in __call__
    self.config_schema = infer_schema_from_config_annotation(
.venv/lib/python3.11/site-packages/dagster/_config/pythonic_config/conversion_utils.py:300: in infer_schema_from_config_annotation
    return infer_schema_from_config_class(model_cls)
.venv/lib/python3.11/site-packages/dagster/_config/pythonic_config/config.py:432: in infer_schema_from_config_class
    fields[resolved_field_name] = _convert_pydantic_field(pydantic_field_info)
.venv/lib/python3.11/site-packages/dagster/_config/pythonic_config/conversion_utils.py:139: in _convert_pydantic_field
    config_type = _config_type_for_type_on_pydantic_field(field_type)
.venv/lib/python3.11/site-packages/dagster/_config/pythonic_config/conversion_utils.py:197: in _config_type_for_type_on_pydantic_field
    return Noneable(_config_type_for_type_on_pydantic_field(optional_inner_type))
.venv/lib/python3.11/site-packages/dagster/_config/pythonic_config/conversion_utils.py:202: in _config_type_for_type_on_pydantic_field
    return Map(
.venv/lib/python3.11/site-packages/dagster/_config/field_utils.py:188: in __init__
    check.inst_param(self.key_type, "key_type", ConfigType)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

obj = False, param_name = 'key_type', ttype = <class 'dagster._config.config_type.ConfigType'>, additional_message = None

    def inst_param(
        obj: T, param_name: str, ttype: TypeOrTupleOfTypes, additional_message: Optional[str] = None
    ) -> T:
        if not isinstance(obj, ttype):
>           raise _param_type_mismatch_exception(
                obj, ttype, param_name, additional_message=additional_message
            )
E           dagster._check.ParameterCheckError: Param "key_type" is not a ConfigType. Got False which is type <class 'bool'>.

.venv/lib/python3.11/site-packages/dagster/_check/__init__.py:640: ParameterCheckError

Deployment type

None

Deployment details

No response

Additional information

No response

Message from the maintainers

Impacted by this issue? Give it a 👍! We factor engagement into prioritization.

[benpankow]

Hi Daniel, will take a look at how difficult this would be to support next week

[rex-asabor]

any updates here? @benpankow