I'd prefer if we didn't use Any here, since this is a publicly exposed attribute

We could resolve it using a TypeVar:
T = TypeVar("T", bound=Union[Expression, CONSTANT_TYPE])

But honestly, I'd recommend converting the constants at init time instead, which would both resolve this issue, raise conversion errors in a better place, and simplify the code:

class Literals(Stage):
"""Returns documents from a fixed set of predefined document objects."""
def __init__(self, *documents: dict[str, Expression | CONSTANT_TYPE]):
super().__init__("literals")
self.documents: Sequence[dict[str, Expression]] = [
{k: v if isinstance(v, Expression) else encode_value(v) for k,v in d.items()}
for d in documents
]
def _pb_args(self):
return [Value(map_value={"fields": doc}) for doc in self.documents]

I'm not sure if it's a good idea because encoding the values here could break the __repr__() method at the parent class, and we would expose the encoded values in the repr string. For example, instead of "Literals(documents=('hello',))", we would be getting "Literals(documents=(Value(string_value='hello'),))".

What do you think?

Sorry, I should have tested the code before posting. Value isn't an expression, so that wouldn't have passed mypy. This should though:

class Literals(Stage):
"""Returns documents from a fixed set of predefined document objects."""
def __init__(self, *documents: Mapping[str, Expression | CONSTANT_TYPE]):
super().__init__("literals")
self.documents: list[Mapping[str, Expression]] = [
{k: v if isinstance(v, Expression) else Constant.of(v) for k,v in d.items()}
for d in documents
]
def _pb_args(self):
return [
Value(map_value={"fields": {k: v.to_pb for k,v in doc.items()}})
for doc in self.documents
]

This would result in the repr of Literals({"hello": "world"})) being Literals(documents=[{"hello": Constant.of('world')}]), which is fine, because represents the same object. And that's how the other stages are set up