Is it possible to create a TypedDict equivalent with its own methods and properties in Python?

Question

I like to work with TypedDicts to convey the batch and item information in PyTorch. It works well with the default collate function and adds the possibility to add types.

What I don't like is that you need to call it with strings keys, that is error prone, and that you can't define methods.

I usually overcome that by embedding the definition of the dicts as well as the keys and methods in a module. For example:

from typing import TypedDict

class Keys:
    source = "source"
    encoding = "encoding"
    prediction = "prediction"

class Batch(TypedDict):
    source: list[str]
    encoding: Tensor  # Shape NCHV
    prediction: Tensor  # Shape N

class Item(TypedDict):
    source: str
    encoding: Tensor  # CHW
    result: float

def item_at(index: int, batch: Batch):
    """Extract an item from the batch"""
    return Item(...)

def get_source(x: Batch | Item):
    return x[Keys.source]

def set_source(x: Batch | Item, value: str | List[str]):
    x[Keys.source] = value

I think it could be more interesting to embed the methods into the Batch and Item classes.

However that's not allowed by TypedDict. On the other hand, when I inherit from dict, I don't know how to add the types.

Do you know how to to do the trick and have a TypedDict equivalent with its own methods and properties?