Python Protocol with __getitem__ and typed keys

  python, python-typing

I am looking for the correct type to annotate this code

def process_item(item: MyItemType) -> str:
    return f"Hello {item['name']} {item['age']} old"

i.e item is any object with __getitem__() and must have name and age keys.

import typing

class MyItemType(typing.Protocol):
    name: str
    age: int

    def __getitem__(self, key: ???)->???:
        ...

In practice it would mean valid inputs include
pyspark.sql.Row instance, vanilla dict or custom object

class MyItem:
    def __init__(self, name:str, age: int):
        self.name = name
        self.age = age
    def __getitem__(self, k: ???) -> ???:
        if not hasattr(self, k):
            raise KeyError()
        return getattr(self, k)

I can do

class MyItemType[typing.TypedDict]:
    name: str
    age: int

def process_item(item: MyItemType):...

but since that is not a typing.Protocol this wouldn’t work passing in the pyspark.sql.Row or custom class with these attributes.

Source: Python Questions

LEAVE A COMMENT