Subsetting multi-assets code snippet in documentation has a bug

[BenameurZyad]

What's the issue or suggestion?

The code snippet provided as an example in the documentation here does not work correctly :
https://docs.dagster.io/concepts/assets/multi-assets#subsetting-multi-assets

Here is the code from the documentation :

from dagster import AssetExecutionContext, AssetSpec, MaterializeResult, multi_asset


@multi_asset(
    specs=[AssetSpec("asset1", skippable=True), AssetSpec("asset2", skippable=True)],
    can_subset=True,
)
def split_actions(context: AssetExecutionContext):
    if "asset1" in context.op_execution_context.selected_asset_keys:
        yield MaterializeResult(asset_key="asset1")
    if "asset2" in context.op_execution_context.selected_asset_keys:
        yield MaterializeResult(asset_key="asset2")

The if conditions are not working correctly.
The context.op_execution_context.selected_asset_keys returns a Set[AssetKey].
A possible fix could be :

# if "asset1" in selected_asset_keys
if AssetKey("asset1") in selected_asset_keys

Here is a debugging version of the code :

from dagster import AssetExecutionContext, AssetKey, AssetSpec, Definitions, MaterializeResult, multi_asset


@multi_asset(
    specs=[AssetSpec("asset1"), AssetSpec("asset2")],
    can_subset=True,
)
def split_actions(context: AssetExecutionContext):

    # returned asset keys are of type set[AssetKey]
    selected_asset_keys = context.op_execution_context.selected_asset_keys
    context.log.info(f"selected_asset_keys: {selected_asset_keys}")
    context.log.info(f"selected_asset_keys type: {type(selected_asset_keys)}")

    if AssetKey("asset1") in selected_asset_keys:
        string_in_selected_asset_keys = "asset1" in selected_asset_keys
        asset_key_in_selected_asset_keys = AssetKey("asset1") in selected_asset_keys
        context.log.info(f"Finding asset name (string) in selected asset keys: {string_in_selected_asset_keys}")
        context.log.info(f"Finding asset key in selected asset keys: {asset_key_in_selected_asset_keys}")
        yield MaterializeResult(asset_key="asset1")

    if AssetKey("asset2") in selected_asset_keys:
        string_in_selected_asset_keys = "asset2" in selected_asset_keys
        asset_key_in_selected_asset_keys = AssetKey("asset2") in selected_asset_keys
        context.log.info(f"Finding asset name (string) in selected asset keys: {string_in_selected_asset_keys}")
        context.log.info(f"Finding asset key in selected asset keys: {asset_key_in_selected_asset_keys}")
        context.log.info("Materializing asset2")
        yield MaterializeResult(asset_key="asset2")

defs = Definitions(assets=[split_actions])

Additional information

No response

Message from the maintainers

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