Quickly locate an item in a dataclass containing a list of dataclasses by field value

Viewed 1335

I have a dataclass with this structure:

from dataclasses import dataclass
from typing import List

@dataclass
class PartData:
    id: int = 0
    name: str = None
    value: int = 0

@dataclass
class StockData:
    stock_1: List[PartData] = None
    stock_2: List[PartData] = None
    def __getitem__(self, key):
        return super().__getattribute__(key)

Now I create the dataclasses and fill them with items:

PARTS = [{"id": 1, "name": "screw"}, {"id": 3, "name": "bolt"}, {"id": 42, "name": "glue"}, {"id": 11, "name": "nail"}, {"id": 31, "name": "hammer"}, {"id": 142, "name": "paper"}]

dc_stock = StockData()

for p in PARTS:
    dc_part = PartData()
    dc_part.id = p["id"]
    if dc_part.id % 2 == 0:
        dc_stock_list = "stock_1"
    else:
        dc_stock_list = "stock_2"
    if getattr(dc_stock, dc_stock_list) == None:
        setattr(dc_stock, dc_stock_list, [dc_part])
    else:
        dc_stock[dc_stock_list].append(dc_part)

print(dc_stock)
# StockData(stock_1=[PartData(id=42, name=None, value=0), PartData(id=142, name=None, value=0)], 
#           stock_2=[PartData(id=1, name=None, value=0), PartData(id=3, name=None, value=0), PartData(id=11, name=None, value=0), PartData(id=31, name=None, value=0)]) 

I know I can loop over all items and compare them, but can I define a method that takes part_id as an argument and can update any item in dc_stock with that part_id with a new value? Can this be implemented as a method of StockData? Suppose I do not know if the part is in stock_1 or stock_2.

Edit

For better understanding I want to share my approach, which looks very loopy and costy to me:

@dataclass
class StockData:
    stock_1: List[PartData] = None
    stock_2: List[PartData] = None

    def __getitem__(self, key):
        return super().__getattribute__(key)

    def update_part(self, id, value):
        for stock_list in [f for f in fields(self) if f.name.startswith("stock")]:
            stock = getattr(self, stock_list.name)
            if len(stock) > 0:
                for part in stock:
                    if part.id == id:
                        part.value = value
                        return None


print(dc_stock)
dc_stock.update_part(1, 10)
print(dc_stock)
2 Answers

Here's one way to set it up. If you always know you need to lookup by id, you can use a dict mapping of id to part instead, since a dict lookup is much faster than finding an part from a list. I also cache the list of dataclass fields that are related to stocks, just as that might be a good idea also.

from dataclasses import dataclass, fields, field
from functools import cached_property
from typing import List, Dict, Union, Tuple


@dataclass
class PartData:
    id: int = 0
    name: str = None
    value: int = 0


@dataclass
class StockData:
    stock_1: Dict[int, PartData] = field(default_factory=dict)
    stock_2: Dict[int, PartData] = field(default_factory=dict)

    @cached_property
    def stock_fields(self) -> Tuple[str, ...]:
        return tuple(f.name for f in fields(self)
                     if f.name.startswith("stock"))

    @classmethod
    def from_parts(cls, parts: List[Dict[str, Union[str, int]]]):
        """Create a new `StockData` object from list of parts."""
        stock = cls()

        for p in parts:
            part = PartData(**p)
            if part.id % 2 == 0:
                stock_list = 'stock_1'
            else:
                stock_list = 'stock_2'

            getattr(stock, stock_list)[part.id] = part

        return stock

    def update_part(self, id, value):
        """Update value for a part, given the part id."""

        for stock_field in self.stock_fields:
            stock = getattr(self, stock_field)
            if id in stock:
                stock[id].value = value
                return None

Usage is pretty similar to how you had it. I also added a from_parts helper method, as it seems it might be a common pattern to construct a StockData instance from a list of parts. Note that since the stock fields are now dictionaries, you can access the .values() to iterate over the PartData items in each stock.

def main():
    PARTS = [{"id": 1, "name": "screw"}, {"id": 3, "name": "bolt"},
             {"id": 42, "name": "glue"}, {"id": 11, "name": "nail"},
             {"id": 31, "name": "hammer"}, {"id": 142, "name": "paper"}]

    dc_stock = StockData.from_parts(PARTS)
    assert dc_stock.stock_2[1].value == 0

    print(dc_stock)

    dc_stock.update_part(1, 10)
    assert dc_stock.stock_2[1].value == 10

    print(dc_stock)

    print('Stock 1:')
    print(list(dc_stock.stock_1.values()))


if __name__ == '__main__':
    main()

What you're asking for is called indexing.

Basically, you have a dict accompanying your data structure of {<field value>: <items with this value>} which is updated appropriately whenever you update the data.
It's even easier if the field is unique (as an item ID should be): you need to only link to 1 item from a key rather than to a list of items.

As you can see, keeping an index up to date is extra work, so it'd only benefit you past a certain data size; it also matters how often the data is written vs read (an index costs time on an update but saves time on a select past a certain data size once index lookup overhead becomes faster than iterating over the entire table) and which percentage of the queries will benefit from the index.


First of all, consider not reinventing the wheel and using a Pythonic ORM like SQLAlchemy instead of dataclasses that does support indexing transparently. You don't need to run a DB server to benefit from it as it can use a serverless DB like SQLite as a backend, too. Moreover, a compiled backend will likely be much (orders of magnitude) faster than a pure Python one.


The way to integrate a dict-base index into your data structure would be to keep it in the table class (StockData) and command the table instance to update the index whenever any of the indexed fields is written (including when they are first initialized).

  • Probably the easiest way to do that is:
    • Keep a reference to the table instance in each record instances (it's sufficient to keep a reference to just the update method)
    • Use it to command the table instance to update indices whenever an indexed field is written (incl. when it's first initialized)
  • If you don't want to modify field classes, your options are:
    • Do not write fields directly but only through some interface provided by the table class. This way, the table class' logic will have an opportunity to update the index because it will get control after writing the value but before returning to you
    • Update the index manually after any writes. This is error-prone (=a recipe for eventual disaster), especially in more complex operations with many interdependent steps 'cuz you may forget, or even not be allowed by syntax (e.g. if you use generator expressions) to call the update at an appropriate moment.

Here's an illustration of what the "easiest way" option can look like:

class PartData:
  <...>
  _table: StockData

  def __setitem__(self, key, new_value):
    if key == 'id':
      self._table.update_id_index(self, new_value, self.__getitem__(key))
    super(self,PartData).__setitem__(self, key, new_value)

class StockData:
  <...>
  # assuming id is unique
  id_index: {object: PartData} = {}

  def update_id_index(self, record, new_value, old_value = None):
    try: del self.id_index[old_value]
    except KeyError: pass
    self.id_index[new_value] = record
Related