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