Generic Types Python

main.py

# Run the type checks installing mypy and executing mypy main.py

import asyncio

from typing import TypeVar, Protocol, Generic

# Any class that implements the _UserPorfile protocol (static duck typing) 
# will be accepted by update_user_profile function as a valid profile
class UserProfile:
    """User profile class."""	
    def __init__(self, first_name: str, last_name: str, birth_date: str, other: int) -> None:
        self.first_name = first_name
        self.last_name = last_name
        self.birth_date = birth_date

    def __repr__(self) -> str:
        return f'<UserProfile({self.first_name}, {self.last_name}, {self.birth_date})>'

class BadUserProfile:
    """Bad user profile class."""
    def __init__(self, first_name: str, last_name: str) -> None:
        self.first_name = first_name
        self.last_name = last_name

    def __repr__(self) -> str:
        return f'<BadUserProfile({self.first_name}, {self.last_name})>'


T = TypeVar('T', bound='_UserProfile')

class _UserProfile(Protocol):
    """User profile protocol."""	
    first_name: str
    last_name: str
    birth_date: str


# This function will be used on different places on the services.
async def update_user_profile(profile: T) -> T:
    """Updates user profile.
    :param profile: User profile.
    :type profile: _UserProfile protocol.
    :return: Updated user profile.
    """
    await asyncio.sleep(1) # Simulate some async work
    
    # Here we can acces to all attributes with type checking and hints
    print(profile.first_name, profile.last_name, profile.birth_date)
    # bad_profile = BadUserProfile(profile.first_name, profile.last_name) 
    # return bad_profile # MyPy error: Incompatible return value type (got "BadUserProfile", expected "T")
    return profile


class UserProfileClient(Generic[T]):
    """User profile client."""
    async def update_user_profile(self, profile: T) -> T:
        """Updates user profile.
        :param profile: User profile.
        :type profile: _UserProfile protocol.
        :return: Updated user profile.
        """
        await asyncio.sleep(1) # Simulate some async work
        
        # Here we can acces to all attributes with type checking and hints
        print(profile.first_name, profile.last_name, profile.birth_date)
        # bad_profile = BadUserProfile(profile.first_name, profile.last_name)
        # return bad_profile # MyPy error: Incompatible return value type (got "BadUserProfile", expected "T")
        return profile


async def main() -> None:
    # Run type checking with MyPy: mypy main.py
    profile = UserProfile('John', 'Doe', '1990-01-01', 1)
    # bad_profile = BadUserProfile('John', 'Doe')
    updated_profile = await update_user_profile(profile) # MyPy Success: no issues found in 1 source file 
    # bad_updated_profile = await update_user_profile(bad_profile) # MyPy error: Value of type variable "T" of "update_user_profile" cannot be "BadUserProfile"
    print(updated_profile)
    # print(bad_updated_profile)

if __name__ == '__main__':
    asyncio.run(main())