4 - Classes and DataClasses

Behavior-focused and data-focused class.

class BankAccount:
    def __init__(self, account_holder: str, balance: int = 0):
        self.account_holder = account_holder
        self.balance = balance

    def deposit(self, amount: int):
        self.balance += amount

    def withdraw(self, amount: int):
        if amount <= self.balance:
            self.balance -= amount
            return True
        else:
            return False

    # Comparison based on account balance
    def __lt__(self, other: "BankAccount") -> bool:
        return self.balance < other.balance

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, BankAccount):
            raise NotImplementedError
        return self.balance == other.balance


def main() -> None:
    account1 = BankAccount("Alice", 1000)
    account2 = BankAccount("Bob", 1500)

    print(account1 < account2)  # True, because Alice's balance is less than Bob's
    print(account1 == account2)  # False, because their balances are not equal


if __name__ == "__main__":
    main()

True
False

Dataclass

In normal behavior-focused class, we need to pass in the variables to the initializer (__init__) and also store that in the instance, which is a duplication and extra code that we dont want to have.

def __init__(self, account_holder: str, balance: int = 0):
    self.account_holder = account_holder
    self.balance = balance

Since Python 3.7 it offers dataclass. For the BankAccount class that we defined above, we can modify it using dataclass like this:

from dataclasses import dataclass

@dataclass
class BankAccount:
    account_holder: str
    balance: int=0

from dataclasses import dataclass
from enum import StrEnum, auto

class FuelType(StrEnum):
    PETROL = auto()
    DIESEL = auto()
    ELECTRIC = auto()

@dataclass
class Vehicle:
    brand: str
    model: str
    color: str
    license_plate: str
    driving_miles: int=0
    fuel_type: FuelType = FuelType.ELECTRIC

    def need_maintenance(self, maintenance_miles: int = 10000) -> bool:
        return self.driving_miles >= maintenance_miles
    
def main() -> None:
    tesla = Vehicle(
        brand="Tesla",
        model="Model 3",
        color="black",
        license_plate="ABC-123"
    )

    print(tesla)

if __name__ == "__main__":
    main()

Vehicle(brand='Tesla', model='Model 3', color='black', license_plate='ABC-123', driving_miles=0, fuel_type=<FuelType.ELECTRIC: 'electric'>)

Dunder methods

Class without dataclass

# Class without dataclass
class Person:
    def __init__(self, name: str, age: int):
        self.name: str = name
        self.age: int = age

    def __repr__(self):
        return f"Person(name={self.name}, age={self.age})"

    def __eq__(self, other: object):
        if not isinstance(other, Person):
            return False
        return self.name == other.name and self.age == other.age

    def __ne__(self, other: object):
        if not isinstance(other, Person):
            return False
        return not self.__eq__(other)

Class with dataclass

@dataclass
class PersonDataclass:
    name: str
    age: int


def main() -> None:
    person1 = Person("Alice", 30)
    person2 = Person("Bob", 25)

    person_dataclass1 = PersonDataclass("Alice", 30)
    person_dataclass2 = PersonDataclass("Bob", 25)

    
    # Check if instances are equal
    print(person1 == person2)  # Output: False
    print(person_dataclass1 == person_dataclass2)  # Output: False

    print(person1 != person2)  # Output: True
    print(person_dataclass1 != person_dataclass2)  # Output: True

    # String representation of instances
    print(repr(person1))  # Output: Person(name=Alice, age=30)
    print(repr(person_dataclass1))  # Output: PersonDataclass(name='Alice', age=30)

if __name__ == "__main__":
    main()

False
False
True
True
Person(name=Alice, age=30)
PersonDataclass(name='Alice', age=30)