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

You need to define `__getitem__` – Asocia Jun 24 '20 at 17:37 — Asocia, Jun 24 '20 at 17:37

Asocia · Accepted Answer · 2022-03-25T11:53:19.237

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

score 3 · Answer 2 · answered Feb 27 '21 at 19:15

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}

How to make custom data class subscriptable

2 Answers2

With pydantic

Without pydantic