The cloud-native orchestrator for the whole development lifecycle, with integrated lineage and observability.

dagster

I’m trying to set up a custom IO manager that handles partition keys, and I’m a bit confused about how to implement `load_input` properly.

In the case where a partitioned asset materialization is loading a non-partitioned asset, it seems like `InputContext.partition_key` is still set to the partition key of the materialization. Does that mean I cannot deduce from the context passed to `load_input` whether the upstream asset is partitioned or not?

After looking at `UPathIOManager` , which is the base class for the `fs_io_manager` , it seems like looking at `context.has_asset_partitions`  and `context.asset_partition_keys` is the way to go:
<https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/dagster/_core/storage/upath_io_manager.py#L134-L205>