41 lines
938 B
Python
41 lines
938 B
Python
from dataclasses import dataclass
|
|
from typing import Generic, List, Optional, TypeVar
|
|
|
|
|
|
@dataclass
|
|
class UserBase:
|
|
id: int
|
|
first_name: str
|
|
last_name: str
|
|
|
|
|
|
@dataclass
|
|
class Employee(UserBase):
|
|
department: str
|
|
|
|
|
|
@dataclass
|
|
class Partner(UserBase):
|
|
company: str
|
|
|
|
|
|
UserContainerType = TypeVar("UserContainerType", bound=UserBase)
|
|
|
|
|
|
class UserContainer(List[UserContainerType]):
|
|
def get_user_by_id(self, id: int) -> Optional[UserContainerType]:
|
|
for user in self:
|
|
if user.id == id:
|
|
return user
|
|
return None
|
|
|
|
|
|
class EmployeeContainer(UserContainer[Employee]):
|
|
def get_by_department(self, department: str) -> List[Employee]:
|
|
return [user for user in self if user.department == department]
|
|
|
|
|
|
class PartnerContainer(UserContainer[Partner]):
|
|
def get_by_company(self, company: str) -> List[Partner]:
|
|
return [user for user in self if user.company == company]
|