Search code examples
pythonpython-dataclassespython-class

Python dataclass, one attribute referencing other


@dataclass
class Stock:
    symbol: str
    price: float = get_price(symbol)

Can a dataclass attribute access to the other one? In the above example, one can create a Stock by providing a symbol and the price. If price is not provided, it defaults to a price which we get from some function get_price. Is there a way to reference symbol?

This example generates error NameError: name 'symbol' is not defined.


Solution

  • You can use __post_init__ here. Because it's going to be called after __init__, you have your attributes already populated so do whatever you want to do there:

    from typing import Optional
    from dataclasses import dataclass
    
    
    def get_price(name):
        # logic to get price by looking at `name`.
        return 1000.0
    
    
    @dataclass
    class Stock:
        symbol: str
        price: Optional[float] = None
    
        def __post_init__(self):
            if self.price is None:
                self.price = get_price(self.symbol)
    
    
    obj1 = Stock("boo", 2000.0)
    obj2 = Stock("boo")
    print(obj1.price)  # 2000.0
    print(obj2.price)  # 1000.0
    

    So if user didn't pass price while instantiating, price is None. So you can check it in __post_init__ and ask it from get_price.


    There is also another shape of the above answer which basically adds nothing more to the existing one. I just added for the records since someone might attempt to do this as well and wonder how is it different with the previous one:

    @dataclass
    class Stock:
        symbol: str
        price: InitVar[Optional[float]] = None
    
        def __post_init__(self, price):
            self.price = get_price(self.symbol) if price is None else price
    

    You mark the price as InitVar and you can get it with a parameter named price in the __post_init__ method.