Create a union of all supported document names and documents to allow type inference

[coretl]

Callbacks to the bluesky RE are currently called as callback(name, doc) which does not allow inferring the structure of doc when given a particular name. If we passed a tuple of the two and called callback(named_doc) then we could do type inference on the content. E.g.

Code sample in pyright playground

from typing import TypedDict, Literal, Unpack

class DocA(TypedDict):
    foo: str
    bar: int

class DocB(TypedDict):
    foo: str
    bat: float

NamedDocument = \
    tuple[Literal["doc_a"], DocA] | \
    tuple[Literal["doc_b"], DocB]


def my_callback(named_document: NamedDocument):
    match named_document:
        case "doc_a", doc:
            num = doc["bar"]
        case "doc_b", doc:
            num = doc["bat"]
    print(num)

This ticket is to decide if this is a good idea, then create the NamedDocument union so that a corresponding bluesky ticket can be created to discuss the callback changes on the bluesky time.

Acceptance Criteria

• Decide this is a good idea, then create the union

[coretl]

@jacopoabramo @evvaaaa @tacaswell @danielballan thoughts?

[jacopoabramo]

Apparently this works with pyright but not with mypy: https://mypy-play.net/?mypy=latest&python=3.10&gist=3f542f6ee43a1455e9ef67d885d9ba37

mypy doesn't seem to be able to narrow down a union of multiple tuples...

[tacaswell]

I think this is a good idea and the minimal change.

We should add a decorator like

def new_signature_decorator(func: Callable[...]) -> Callable[...]:
    def inner(named_document: NamedDocument):
        return func(*named_document)

    # metadata patching
    return inner

in event model to make it easy to migrate to the new signature if we want to make RE warn/error on the old signature.

[coretl]

@jacopoabramo do you need mypy?

[jacopoabramo]

@coretl it's my go-to solution for type checking, but if pyright is the preferred approach of the community and you see this as a valuable addition I'll make my peace with it so don't feel blocked by me

[coretl]

If we support mypy then we have to make an extra dict wrapper layer:
https://mypy-play.net/?mypy=latest&python=3.14&gist=99529f427d0527c08fd8518c575ac642

Worth it?

[jacopoabramo]

We're sort of coming back to using tagged unions but at least it's compatible with what I had in mind in the beginning so I'm happy with that

[coretl]

Ok, this works with pyright too, it's uglier, but is it worth it to support both mypy and pyright?

Code sample in pyright playground

from typing import TypedDict, Literal, TypeVar, Generic

class DocA(TypedDict):
    foo: str
    bar: int

class DocB(TypedDict):
    foo: str
    bat: float

NameT = TypeVar("NameT")
DocT = TypeVar("DocT")

class DocWrapper(TypedDict, Generic[NameT, DocT]):
    name: NameT
    doc: DocT
    

NamedDocument = \
    DocWrapper[Literal["doc_a"], DocA] | \
    DocWrapper[Literal["doc_b"], DocB]


def my_callback(named_document: NamedDocument):
    match named_document["name"]:
        case "doc_a":
            num = named_document["doc"]["bar"]
        case "doc_b":
            num = named_document["doc"]["bat"]
    print(num)

[coretl]

@evvaaaa thoughts?

[coretl]

Ok, the best I can do with mypy is a NamedTuple, but you still need to mention the class name in the case statement, otherwise it falls back to the tuple form.

mypy

Code sample in mypy playground

from typing import TypedDict, Literal, TypeVar, Generic, NamedTuple

class DocA(TypedDict):
    foo: str
    bar: int

class DocB(TypedDict):
    foo: str
    bat: float

NameT = TypeVar("NameT")
DocT = TypeVar("DocT")

class NamedDocument(NamedTuple, Generic[NameT, DocT]):
    name: NameT
    doc: DocT
    

NamedDocumentUnion = \
    NamedDocument[Literal["doc_a"], DocA] | \
    NamedDocument[Literal["doc_b"], DocB]


def my_callback(named_document: NamedDocumentUnion):
    num: int | float
    match named_document:
        case NamedDocument("doc_a", doc):
            num = doc["bar"]
        case NamedDocument("doc_b", doc):
            num = doc["bat"]
    print(num)

pyright

Code sample in pyright playground

from typing import TypedDict, Literal, TypeVar, Generic, NamedTuple, reveal_type

class DocA(TypedDict):
    foo: str
    bar: int

class DocB(TypedDict):
    foo: str
    bat: float

NameT = TypeVar("NameT")
DocT = TypeVar("DocT")

class NamedDocument(NamedTuple, Generic[NameT, DocT]):
    name: NameT
    doc: DocT
    

NamedDocumentUnion = \
    NamedDocument[Literal["doc_a"], DocA] | \
    NamedDocument[Literal["doc_b"], DocB]


def my_callback(named_document: NamedDocumentUnion):
    match named_document:
        case "doc_a", doc:
            num = doc["bar"]
        case "doc_b", doc:
            num = doc["bat"]
    reveal_type(num)

[jacopoabramo]

I don't mind this actually! It's very explicit

[evvaaaa]

I don't mind too much either... just more cumbersome than being able to unpack.

I worry that python ends up adding some nicer runtime checkable type discriminator support for TypedDict and we make it even more difficult than what we have already in regards to backwards compatibility if we want to support it.

I think I'd still vote for adding a document_type field to each document because

1. A consumer already shouldn't be verifying an exact match to the schema + exclusively their own consumer for a run_start. A consumer probably shouldn't be assuming that every new key in the top level of the run start is for them and raising an error if it's not recognized.
2. I don't see the likelyhood of many consumers having document_type as a top level argument, unless they have some project where they have created further types of run start for use in their consumer, in which case they should probably be using "MY_CONSUMER_NAME": dict instead.

I think a small minority of projects having to make a patch release for this would be worth the future proofing.

[coretl]

I'm ambivalent as to whether we add document_type to the documents, or invent a NamedDocument namedtuple wrapper like the last example.

I'll let @tacaswell or @danielballan make the decision...

[tacaswell]

I remain in favor of NamedDocument namedtuple wrapper. It is the minimal change over the status quo and keeps routing level information (what is this thing and who wants it?) in a different band then the data (doing something useful with this thing).

Treating it as a (name, doc) pair is more-or-less exactly what we already do in the zmq publisher:

https://github.com/bluesky/bluesky/blob/3f23fd610bb5cc7fd90fd3ed58e6a04544295c47/src/bluesky/callbacks/zmq.py#L89-L92

---

Way back in the day we used to have a dictionary sub-class (called Doct) that added the document name as an attribute (and also provided .name access to the values and tried to be read-only) that we abandoned for being more trouble than it was worth.