Type mismatch for `context.partition_key`

[drawnwren]

Dagster version

1.2.7

What's the issue?

InputContext and OutputContext .partition_key has type -> str but then, in a MultiPartitionDefinition asset -- the docs recommend that we do context.partition_key.keys_by_dimension[dim.name] . Pyright really doesn't like this. I'm not sure if this is a bug or if it's just my lsp being annoying.

What did you expect to happen?

No response

How to reproduce?

No response

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.

[yuhan]

@clairelin135 mind taking a look?

[clairelin135]

Hey @DrawnWren. Updating the type handling for these context methods is something we wanted to add, but unfortunately it would constitute a breaking change for our users who expect the methods to return a string.

I think this is something we can revisit close to a major release--will just keep this in the backlog until then.

[sryza]

@DrawnWren - the background here is that MultiPartitionKey is a subclass of string. context.partition_key can't have a MultiPartitionKey return type annotation, because, when not using multi-partitions, it won't return this type.

@clairelin135 I think there area a couple solutions we could consider here:

• Update our examples to include a cast: cast(MultiPartitionKey, context.partition_key). keys_by_dimension[dim_name]
• Add a method to contexts named something like get_partition_key_dimension that accepts a dimension name.

[clairelin135]

Add a method to contexts named something like get_partition_key_dimension that accepts a dimension name.

This option sounds great to me, I'll go ahead and add this