How to make custom data class subscriptable

Question

Consider this data class derived from the pydantic package:

from typing import List
from pydantic import BaseModel 
class Bucket(BaseModel):
    setting: List[str]
    fight_1: List[int]
    cause_1: List[str]

let my_bucket be an instance of Bucket:

my_bucket = Bucket(setting=['some_value'], fight_1=[0], cause_1=['other_value'])

basically I would like to be able to do

my_bucket['setting']                                                                                                         

and get back ['some_value'], but instead I get:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-18-cacbdc1698e9> in <module>
----> 1 my_bucket['setting']

TypeError: 'Bucket' object is not subscriptable

Answer 1

With pydantic

You can access your property as an attribute in pydantic.

my_bucket = Bucket(setting=['some_value'], fight_1=[0], cause_1=['other_value'])
print(my_bucket.setting)  # ['some_value']

If you want to use [] to access it, you need to define __getitem__:

from typing import List
from pydantic import BaseModel


class Bucket(BaseModel):
    setting: List[str]
    fight_1: List[int]
    cause_1: List[str]

    def __getitem__(self, item):
        return getattr(self, item)


my_bucket = Bucket(setting=['some_value'], fight_1=[0], cause_1=['other_value'])
print(my_bucket['setting'])  # ['some_value']

Without pydantic

If you don't want to use pydantic and create your custom dataclass you can do this:

from dataclasses import dataclass


@dataclass
class CustomDataClass:
    data: int

    def __getitem__(self, item):
        return getattr(self, item)


obj = CustomDataClass(42)
print(obj.data)  # 42
print(obj["data"])  # 42, needs __getitem__ to be implemented

Answer 2

For a python dataclass this works for me:

from dataclasses import dataclass

@dataclass
class MyCustomRequest:
    request: Dict
        
    def __getitem__(self, key):
        return super().__getattribute__(key)

>>> a = MyCustomRequest({"a":10})
>>> a["request"]
{"a":10}
>>> a.request
{"a": 10}