Inheriting Pydantic Models

Inheritance is a fundamental concept in object-oriented programming that allows developers to create new classes based on existing ones, reusing and extending functionality. In this chapter, I will explore how to inherit Pydantic models to create more efficient and maintainable code.

Creating a Base Model

To start, we need to create a base model that serves as the foundation for our inherited models. In Pydantic, a model is defined using the BaseModel class. Let's create a simple example:

from pydantic import BaseModel, EmailStr

class User(BaseModel):
    id: int
    name: str
    email: EmailStr

Extending the Base Model

Now that we have our base User model, we can create a new model that inherits from it. Suppose we want to create a Student model with additional fields, such as enrollment_date and major. We can achieve this by extending the User model as follows:

from datetime import date
from typing import Optional

class Student(User):
    enrollment_date: date
    major: Optional[str] = None

Adding Validators to the Inherited Model

Validators are functions that validate and preprocess data for Pydantic models. They can be added to inherited models to enforce custom rules or modify data before it's stored in the model. To add a validator to our Student model, we use the @validator decorator:

from pydantic import validator

class Student(User):
    enrollment_date: date
    major: Optional[str] = None

    @validator("enrollment_date")
    def validate_enrollment_date(cls, enrollment_date):
        if enrollment_date > date.today():
            raise ValueError("Enrollment date must be in the past")
        return enrollment_date

Overriding Fields and Methods

In some cases, you might want to override a field or method in the inherited model to change its behavior. To do this, simply redefine the field or method in the derived class:

class Staff(User):
    role: str
    email: str  # Overriding the EmailStr type in the User model

    def get_role(self):
        return f"Staff role: {self.role}"

In this example, we overrode the email field to accept any string instead of validating it as an email address. Additionally, we added a new method get_role to the Staff class.

Practical Examples of Pydantic Inheritance

In this chapter, I will delve into practical examples of Pydantic inheritance that demonstrate the power and flexibility of this feature in real-world applications.

User Authentication

Consider a web application that has multiple user types, such as Admin, Customer, and ContentCreator. Each user type may have its own unique set of attributes and functionality. By leveraging Pydantic inheritance, we can create a base User model and then extend it for each user type.

# Base User model
class User(BaseModel):
    id: int
    username: str
    password: str
    email: EmailStr

# Admin model
class Admin(User):
    access_level: int

    def grant_permission(self, target_user: User):
        # Implementation to grant permission to target_user

# Customer model
class Customer(User):
    shipping_address: str
    billing_address: str

    def place_order(self, order: Order):
        # Implementation to place an order

# ContentCreator model
class ContentCreator(User):
    bio: Optional[str] = None

    def create_content(self, content: Content):
        # Implementation to create content

Product Catalog

In an e-commerce application, you might have different types of products, such as Electronics, Clothing, and Furniture. Each product type could have unique attributes. Here's an example of how to use Pydantic inheritance to model these product types:

# Base Product model
class Product(BaseModel):
    id: int
    name: str
    price: float
    description: str

# Electronics model
class Electronics(Product):
    brand: str
    warranty_period: int

# Clothing model
class Clothing(Product):
    size: str
    material: str
    gender: str

# Furniture model
class Furniture(Product):
    dimensions: Tuple[float, float, float]
    weight: float
    material: str

Blog Post Models

In a blogging platform, you might have various types of content, such as Article, Video, and ImageGallery. Each content type could have unique attributes and functionality. You can use Pydantic inheritance to create models for these content types:

# Base Content model
class Content(BaseModel):
    id: int
    title: str
    author: User
    created_at: datetime

# Article model
class Article(Content):
    body: str
    tags: List[str]

# Video model
class Video(Content):
    video_url: HttpUrl
    duration: int

# ImageGallery model
class ImageGallery(Content):
    images: List[HttpUrl]
    thumbnail: HttpUrl

References

https://docs.pydantic.dev/usage/models/