0
\$\begingroup\$

I have the following employee/company system with dataclasses Employee and Company. There is also a LeaveApplication, Hirer and EmployeePromoter that are in a composition based model with company to fulfil functions. Lastly, there are two Enum called CompanyRoles and CompanyTeams.

I need advise on the following:

  • Whether LeaveApplication really needs to be a separate class as it gets messy
  • If I should create methods give_leaves_employee and hire_employee in company as other methods are in Company. I can later reference them from the hirer.
  • If there is any cleaner way to replicate my promotion system without two lists.

I feel it my code gets messy here with the composition model with functions like promote_employee() in Company and having a separate class, accept_resign() and fire_employee() only in company, and hire_employee() and give_leaves() only in composition based classes. Would there be a better and more readable way to structure. I plan to later implement all if those code into a command line interface in a separate file, where all the classes will be imported

from dataclasses import dataclass, field
from enum import Enum, auto


class CompanyRoles(Enum):
    MANAGER = auto()
    LEAD = auto()
    WORKER = auto()
    INTERN = auto()


class CompanyTeams(Enum):
    MARKETING = auto()
    PRODUCT = auto()
    DESIGN = auto()
    SALES = auto()
    FINANCE = auto()

class EmployeePromoter:
    def __init__(self):
        self.employee_hierarchy = [
            CompanyRoles.INTERN,
            CompanyRoles.WORKER,
            CompanyRoles.LEAD,
            CompanyRoles.MANAGER
        ]

        self.role_salary_increment = [
            0,  
            5000,
            15000,
            25000
        ]
    
    def promote_employee(self, employee):
        employee_role_idx = self.employee_hierarchy.index(employee.role)
        employee.role = self.employee_hierarchy[employee_role_idx + 1]
        employee.salary_dollars += self.role_salary_increment[employee_role_idx + 1]
        print(f"{employee.name} was promoted to a package of {employee.salary_dollars}.")

class Hirer:
    def __init__(self, company) -> None:
        self.company = company

    def employee_eligible(self, employee, role, team):
        employee_has_work_experience = employee.work_experience_years >= self.company.minimum_work_experience
        employee_has_good_grades = employee.grades_percent_average >= self.company.minimum_hiring_grades
        if self.company.role_occupied(role, team):
            return False
        elif not employee_has_work_experience and not employee_has_good_grades:
            return False
        return True

    def hire_employee(self, applicant, role, team) -> bool:
        if not self.employee_eligible(applicant, role, team):
            print("You did not meet our requirements")
            return False
        self.company.add_employee(applicant)
        self.company.employees = [employee.name for employee in self.company.employee_base]  
        return True    

class LeaveApplication:
    def __init__(self, company) -> None:
        self.company = company

    def give_leaves_employee(self, employee, leaves_required) -> bool:
        ineligible_leaves = leaves_required <= 3 and employee.available_leaves - leaves_required <= 0
        if ineligible_leaves:
            print("Leaves can't be granted.")
            return False
        print("Leaves are granted.")
        return True
    

@dataclass
class Company():

    name : str
    minimum_hiring_grades : int
    minimum_work_experience : int
    employee_leaves : int
    employee_bonus_percent : int
    employee_working_days : int
    employees : list = field(init = False, default_factory = list)
    employee_base : list = field(init = False, default_factory = list)


    def __post_init__(self) -> None:
        '''Creates Hirer object as it can't be created without `self`'''
        self.promoter : object = field(init = False, default = EmployeePromoter())
        self.leave_app : object = field(init = False, default = LeaveApplication(self))
        self.hirer : object = Hirer(self)

    #Database Functions
    def add_employee(self, employee) -> None:
        self.employee_base.append(employee)

    def delete_employee(self, employee) -> None:
        self.employee_base.remove(employee)

    def role_occupied(self, role, team) -> bool:
        return bool([employee.role for employee in self.employee_base if role == employee.role and team == employee.team])

    def promote_employee(self, employee) -> None:
        self.promoter.promote_employee(employee)

    def fire_employee(self, employee) -> None:
        if employee in self.employee_base:
            employee.self_delete_attr()
            self.employee_base.remove(employee)
            self.employees = [employee.name for employee in self.employee_base]  
            print(f"{employee.name} is fired from {self.name}")
        else:
            print(f"{employee.name} is not in the company.")

    def accept_resignation(self, employee) -> None:
        self.employee_base.remove(employee)
        self.employees = [employee.name for employee in self.employee_base]  
        print(f"{employee.name} has now left the company")

    def print_self_employees(self) -> None:
        print(self.employees)




@dataclass
class Employee:
    name : str
    grades_percent_average : int
    work_experience_years : int
    role : Enum = field(init = False, default = None)
    team : Enum = field(init = False, default = None)
    company : object = field(init = False, default= None)
    working_days : int = field(init = False, default= 0)
    bonus_percent : int = field(init = False, default= 0)
    has_job : bool = field(init = False, default= False)
    id : object = field(init = False, default= None)


    def apply_in_company(self, company, role, team) -> None:
        accepted = company.hirer.hire_employee(self, role, team)
        if accepted:
            self.company = company
            self.has_job = True
            self.working_days = company.employee_working_days
            self.bonus_percent = company.employee_bonus_percent
            self.available_leaves = company.employee_leaves
            self.salary_dollars = (self.grades_percent_average + self.work_experience_years) * 1000
            self.role = role
            self.team = team
            self.id = self
    
    def ask_for_leaves(self, leaves_required ) -> None:
        if self.has_job:
            leaves_granted = self.company.leave_app.give_leaves_employee(self, leaves_required) 
            if leaves_granted:
                self.available_leaves -= leaves_required
                self.working_days -= leaves_required
        else:
            print("You do not have a job!")


    def resign_company(self) -> None:
        if self.has_job:
            self.company.accept_resignation(self)
            self.self_delete_attr()
             
        else:
            print("You do not have a job!")

    def self_delete_attr(self) -> None:
        self.company = None
        self.salary_dollars = None
        self.id = None
        self.has_job = False


b = Company(name ="b", minimum_hiring_grades = 90, minimum_work_experience = 5, employee_leaves = 30, employee_bonus_percent = 25, employee_working_days = 300)
arham = Employee(name = "Arham", grades_percent_average = 92, work_experience_years = 5)
rajeev = Employee(name = "Rajeev", grades_percent_average = 91, work_experience_years = 5)
arham.apply_in_company(b, CompanyRoles.WORKER, CompanyTeams.FINANCE)
rajeev.apply_in_company(b, CompanyRoles.WORKER, CompanyTeams.DESIGN)
arham.resign_company()
b.print_self_employees()
```
\$\endgroup\$
4
  • \$\begingroup\$ Are "leaves" vacation days or other paid time off? Is your logic correct there? If leaves_required is 5 and the employee.available_leaves is 0 then it's granted by your code? Also that class doesn't appear to do anything at all with company, so I'm not sure why it needs one for construction or why it's not just a free function. \$\endgroup\$
    – Chuck
    Nov 8, 2022 at 17:52
  • \$\begingroup\$ They are limited vacation days that are not paid time, that every employee receives. If leaves_required is is 5 and employee.available_leaves is 0, leaves will not be granted \$\endgroup\$ Nov 9, 2022 at 10:09
  • \$\begingroup\$ ineligible_leaves = leaves_required <= 3 and employee.available_leaves - leaves_required <= 0 If leaves_required is 5, then leaves_required <= 3 evaluates as False. False and True evaluates to False. This means ineligible_leaves evaluates to False, which means you print("Leaves are granted.") and then return True. \$\endgroup\$
    – Chuck
    Nov 9, 2022 at 15:52
  • \$\begingroup\$ Oh ok, I didn't notice that error, I made that when switching to guard clauses I guess. \$\endgroup\$ Nov 10, 2022 at 10:43

1 Answer 1

1
\$\begingroup\$

LeaveApplication

This doesn't need to be a class:

class LeaveApplication:
    def __init__(self, company) -> None:
        self.company = company

    def give_leaves_employee(self, employee, leaves_required) -> bool:
        ineligible_leaves = leaves_required <= 3 and employee.available_leaves - leaves_required <= 0
        if ineligible_leaves:
            print("Leaves can't be granted.")
            return False
        print("Leaves are granted.")
        return True

Why? self isn't used in its only method. So give_leaves_employee is really a staticmethod:

class LeaveApplication:
    def __init__(self, company) -> None:
        self.company = company

    @staticmethod
    def give_leaves_employee(employee, leaves_required) -> bool:
        ineligible_leaves = leaves_required <= 3 and employee.available_leaves - leaves_required <= 0
        if ineligible_leaves:
            print("Leaves can't be granted.")
            return False
        print("Leaves are granted.")
        return True

Now, since give_leaves is only a staticmethod, let's just make it a standalone function:

def give_leaves_employee(employee, leaves_required) -> bool:
    ineligible_leaves = leaves_required <= 3 and employee.available_leaves - leaves_required <= 0
    if ineligible_leaves:
        print("Leaves can't be granted.")
        return False
    print("Leaves are granted.")
    return True

I also don't love the name. give_leaves_employee doesn't really describe what's happening. Let's call it is_leave_granted, which implies a boolean return:

def is_leave_granted(employee, leaves_required) -> bool:
    ineligible_leaves = leaves_required <= 3 and employee.available_leaves - leaves_required <= 0
    if ineligible_leaves:
        return False

    return True

Let's also break up the boolean condition a bit further:

def is_leave_granted(employee: Employee, leaves_required: int) -> bool:
    too_few_required = leaves_required <= 3 
    too_few_remaining = (employee.available_leaves - leaves_required) <= 0

    if too_few_required and too_few_remaining:
        return False

    return True

Employee

This class probably should be reworked. If we were to be pedantic, an Employee that isn't employed is not really an Employee. They might be considered an Applicant, or more generally, a Person.

I think an Employee should be the output of some hired method on a Company, this way you aren't explicitly looking up Company attributes within the Employee class

@dataclass
class Company:
    ~snip~
    def hire_employee(self, person: Person, role: CompanyRole, team: CompanyTeam) -> Employee:
        salary = (person.grades_percent_average + person.work_experience_years) * 1000
        
        employee = Employee(
            company=self,
            role=role,
            team=team,
            salary=salary,
            **person
        )
        
        self.add_employee(employee)
        return employee



@dataclass
class Employee:
    # Let's take away the is_hired method, since we are putting that
    # on the `Company` class
    def ask_for_leaves(self, leaves_required ) -> None:
        if self.has_job:
            leaves_granted = self.company.leave_app.give_leaves_employee(self, leaves_required) 
            if leaves_granted:
                self.available_leaves -= leaves_required
                self.working_days -= leaves_required
        else:
            print("You do not have a job!")


    def resign_company(self) -> None:
        if self.has_job:
            self.company.accept_resignation(self)
            self.self_delete_attr()
             
        else:
            print("You do not have a job!")

    def self_delete_attr(self) -> None:
        self.company = None
        self.salary_dollars = None
        self.id = None
        self.has_job = False

Now, for this to work, I think it would be nice for this class to be a mapping. This means we can lean on **kwargs style unpacking:

from dataclasses import dataclass, fields

@dataclass
class Employee:
    ~snip~
    
    def keys(self):
        return [field.name for field in fields(self)]
    
    def __getitem__(self, key):
        return getattr(self, key)

    
    ~snip~

Last, I don't think a del_attrs method is necessary. This only happens if an Employee resigns. In this case, I believe that resign should return a Person, since they are unemployed even for the most brief moment

from dataclasses import dataclass, fields

@dataclass
class Person:
    name : str
    grades_percent_average : int
    work_experience_years : int
        
    def keys(self):
        return [field.name for field in fields(self)]
    
    def __getitem__(self, key):
        return getattr(self, key)
    
    

@dataclass
class Employee(Person):
    company: str
    role: CompanyRoles
    team: CompanyTeams
    working_days: int = field(init=False)
    leave_days: int = field(init=False)
    id: object = field(init=False, default=None)

    def __post_init__(self):
        self.working_days = self.company.employee_working_days
        self.leave_days = self.company.employee_leave_days
        self.id = hash(self.name)
        

    def resign(self) -> Person:
        self.company.accept_resignation(self)
        person_attrs = dict(**self)
        
        for k in ('company', 'id', 'role', 'team', 'working_days', 'leave_days'):
            person_attrs.pop(k)
            
        return Person(**person_attrs)
        

c = Company(
    name="AcmeCo", 
    minimum_hiring_grades=90, 
    minimum_work_experience=5, 
    employee_leaves = 30, 
    employee_bonus_percent = 25, 
    employee_working_days = 300
)
p = Person('mike', 4.0, 5)
e = Employee(company=c, role='SomeRole', team='SomeTeam', id='hello', **p)

dict(e)
{'name': 'mike',
 'grades_percent_average': 4.0,
 'work_experience_years': 5,
 'company': 'AcmeCo',
 'id': None}

Now, because we have a Person and an Employee class, we no longer need to check self.employed when is_leave_granted is called.

def is_leave_granted(employee: Employee, leaves_required: int) -> bool:
    too_few_required = leaves_required <= 3 
    too_few_remaining = (employee.available_leaves - leaves_required) <= 0

    if too_few_required and too_few_remaining:
        return False

    return True


@dataclass
class Employee(Person):
    ~snip~
        
    def request_leave(self, leaves_required: int) -> None:
        if not is_leave_granted(self, leaves_required):
            return
        
        self.working_days -= leaves_required
        self.leave_days -= leaves_required
        

Taking a second look, is_leave_granted really is a method on the Employee class:

@dataclass
class Employee(Person):
    def request_leave(self, leaves_required: int) -> None:
        if not self.is_leave_granted(leaves_required):
            return
        
        self.working_days -= leaves_required
        self.leave_days -= leaves_required
        
    def is_leave_granted(self, leaves_required: int) -> bool:
        too_few_required = leaves_required <= 3 
        too_few_remaining = (employee.available_leaves - leaves_required) <= 0

        if too_few_required and too_few_remaining:
            return False

        return True

But it's just wrapping a simple boolean, so it might be better as:

@dataclass
class Employee(Person):
    def request_leave(self, leaves_required: int) -> None:
        too_few_required = leaves_required <= 3 
        too_few_remaining = (employee.available_leaves - leaves_required) <= 0
        
        if too_few_required and too_few_remaining:
            return
        
        self.working_days -= leaves_required
        self.leave_days -= leaves_required
        

Hirer

This factory, in my opinion, is unnecessary and represents a method on the Company class instead:

class Hirer:
    def __init__(self, company) -> None:
        self.company = company

    def employee_eligible(self, employee, role, team):
        employee_has_work_experience = employee.work_experience_years >= self.company.minimum_work_experience
        employee_has_good_grades = employee.grades_percent_average >= self.company.minimum_hiring_grades
        if self.company.role_occupied(role, team):
            return False
        elif not employee_has_work_experience and not employee_has_good_grades:
            return False
        return True

    def hire_employee(self, applicant, role, team) -> bool:
        if not self.employee_eligible(applicant, role, team):
            print("You did not meet our requirements")
            return False
        self.company.add_employee(applicant)
        self.company.employees = [employee.name for employee in self.company.employee_base]  
        return True  

hire_employee is just a wrapper around two things:

  1. Checking eligibility
  2. company.add_employee

So let's fold those into the Company class instead:

@dataclass
class Company():
    name : str
    minimum_hiring_grades : int
    minimum_work_experience : int
    employee_leaves : int
    employee_bonus_percent : int
    employee_working_days : int
    employees : list = field(init = False, default_factory = list)
    employee_base : list = field(init = False, default_factory = list)
    
    ~snip~
    def is_applicant_eligible(self, applicant: Person, role: CompanyRole, team: CompanyTeam) -> bool:
        employee_has_work_experience = applicant.work_experience_years >= self.minimum_work_experience
        employee_has_good_grades = applicant.grades_percent_average >= self.minimum_hiring_grades
        
        if self.role_occupied(role, team):
            return False
        elif not employee_has_work_experience and not employee_has_good_grades:
            return False
        else:
            return True
        
    # we've added this already, but just to show it again
    def hire_employee(self, person: Person, role: CompanyRole, team: CompanyTeam) -> Employee:
        salary = (person.grades_percent_average + person.work_experience_years) * 1000
        
        employee = Employee(
            company=self,
            role=role,
            team=team,
            salary=salary,
            **person
        )
        
        self.add_employee(employee)
        return employee
        

Now, the last thing we need is an apply_for_role method that checks the eligibility and allows a company to hire said employee if eligible

apply_for_role should be a method on the Person class, IMO:

@dataclass
class Person:
    name : str
    grades_percent_average : int
    work_experience_years : int
        
    def keys(self):
        return [field.name for field in fields(self)]
    
    def __getitem__(self, key):
        return getattr(self, key)
    
    def apply_for_role(self, company: Company, role: CompanyRole, team: CompanyTeam) -> Person:
        if not company.is_applicant_eligible(self, role=role, team=team):
            return self
        elif company.role_occupied(role=role, team=team):
            return self
        else:
            return company.hire_employee(self, role, team)

Employees as two lists

You are adding comments that imply that these list.append functions are mocking a database. If this is the case, I'd use a data structure that functions like a database. The best candidate to start with would be a dict:

@dataclass
class Company():
    name : str
    minimum_hiring_grades : int
    minimum_work_experience : int
    employee_leaves : int
    employee_bonus_percent : int
    employee_working_days : int
    employees : dict = field(init = False, default_factory = dict)
        
        
    def add_employee(self, employee: Employee):
        self.employees[employee.name] = Employee
        
        
    def remove_employee(self, employee: Employee):
        # no need to iterate over the whole list
        # when removing employees
        self.employees.pop(employee.name)

    def role_occupied(self, role, team) -> bool:
        # this lookup scales with the number of employees
        it = (
            employee.role for employee in self.employees 
            if role == employee.role and team == employee.team
        )
        return bool(next(filter(bool, it), None))

    def promote_employee(self, employee) -> None:
        self.promoter.promote_employee(employee)
        
    ~snip~

However, it might be worthwhile having another lookup that organizes teams and roles with the optional employee that occupies that role:

class CompanyTeams(Enum):
    MARKETING = auto()
    PRODUCT = auto()
    DESIGN = auto()
    SALES = auto()
    FINANCE = auto()
    
    
@dataclass
class Company():
    name : str
    minimum_hiring_grades : int
    minimum_work_experience : int
    employee_leaves : int
    employee_bonus_percent : int
    employee_working_days : int
    employees : dict = field(init = False, default_factory = dict)
    teams: dict = field(init=False)
        
    def __post_init__(self):
        self.teams = {
            name: {
                role_name: None 
                for role_name in CompanyRoles.__members__
            } for name in CompanyTeams.__members__
        }
        
    def add_employee(self, employee: Employee) -> None:
        self.employees[employee.name] = employee
        self.teams[employee.team.name][employee.role.name] = employee
        
    
    def remove_employee(self, employee: Employee) -> None:
        self.employees.pop(employee.name)
        self.teams[employee.team.name][employee.role.name] = None
        
        
     def role_occupied(self, role: CompanyRole, team: CompanyTeam) -> bool:
        # Now this is a constant time lookup
        return bool(self.teams.get(team.name, {}).get(role.name))
    
    ~snip~

Now, both the lookup, addition, and removal of employees is constant in time.

Teams and Roles

I think the names of these enumerations are too verbose. Call them Team and Role.

Firing Employees

This should return a Person class:

@dataclass
class Company():
    ~snip~
    def fire_employee(self, employee: Employee) -> Person:
        if employee.name not in self.employees:
            raise ValueError(f"Employee {employee.name} not found in company")
            
        self.remove_employee(employee)
        
        attrs = dict(employee)
        
        for k in ('company', 'id', 'working_days', 'leave_days', 'team', 'role'):
            attrs.pop(k)
            
        return Person(**attrs)
    

for k in ...

You'll notice I've used a for k in ('company', ...) to quickly remove attributes from Employee that aren't in Person. This can be refactored by using a global list:

def get_extra_fields():
    person_attrs = {field.name for field in fields(Person)}
    employee_attrs = {field.name for field in fields(Employee)}
    
    return list(employee_attrs - person_attrs)

EXTRA_FIELDS = get_extra_fields()

Now, we can just do:

@dataclass
class Company():
    ~snip~
    def fire_employee(self, employee: Employee) -> Person:
        if employee.name not in self.employees:
            raise ValueError(f"Employee {employee.name} not found in company")
            
        self.remove_employee(employee)
        
        attrs = dict(employee)
        
        for k in EXTRA_FIELDS:
            attrs.pop(k)
            
        return Person(**attrs)

And

@dataclass
class Employee:
    ~snip~
    def resign(self) -> Person:
        self.company.accept_resignation(self)
        person_attrs = dict(**self)
        
        for k in EXTRA_FIELDS:
            person_attrs.pop(k)
            
        return Person(**person_attrs)

Employee.id

This should be assigned by the Company. A naive implementation might be to implement a counter on each Company that forms the id when an employee is hired:

@dataclass
class Company:
    counter: int = 1
        
    def create_id(self):
        """Creates an id like ac00001 for the first employee at AcmeCo"""
        abbrev = self.name[:2].lower()
        num = str(self.counter).zfill(5)
        self.counter += 1
        return ''.join((abbrev, num))
        
        
    def hire_employee(self, person: Person, role: CompanyRole, team: CompanyTeam) -> Employee:
        salary = (person.grades_percent_average + person.work_experience_years) * 1000
        
        employee = Employee(
            company=self,
            role=role,
            team=team,
            salary=salary,
            id=self.create_id(),
            **person
        )
        
        self.add_employee(employee)
        return employee

EmployeePromoter

This also doesn't need to be a class. I'd actually argue that it is two things: 1. A method on the Company class, again 2. A refactor of the Role enum

Why? Because you are effectively re-implementing the Role enum, and it is managing the role of an employee:

# let's change this to set the salary increase for the value
class Role(Enum):
    MANAGER = 25000
    LEAD = 15000
    WORKER = 5000
    INTERN = 0

@dataclass
class Company:
    def promote_employee(self, employee):
        employee_role_idx = self.employee_hierarchy.index(employee.role)
        employee.role = self.employee_hierarchy[employee_role_idx + 1]
        employee.salary_dollars += self.role_salary_increment[employee_role_idx + 1]
        print(f"{employee.name} was promoted to a package of {employee.salary_dollars}.")

But wait, how are you going to work with the index to step the employee forward? Enum is actually built on OrderedDict, and you can use this behavior for consistent iteration:

class Role(Enum):
    MANAGER = 25000
    LEAD = 15000
    WORKER = 5000
    INTERN = 0
    
list(Role)
[<Role.MANAGER: 25000>, <Role.LEAD: 15000>, <Role.WORKER: 5000>, <Role.INTERN: 0>]

# changing the order
class Role(Enum):
    INTERN = 0
    WORKER = 5000
    LEAD = 15000
    MANAGER = 25000
    
list(Role)
[<Role.INTERN: 0>, <Role.WORKER: 5000>, <Role.LEAD: 15000>, <Role.MANAGER: 25000>]

Leaning on this:

@dataclass
class Company:
    def promote_employee(self, employee: Employee):
        # The next role is now definitively next in the iteration
        next_role = next((role for role in Role if role.value > employee.role.value), None)
        
        # What if the employee is a manager?
        if next_role is None:
            raise Exception("Can't promote a manager!")
            
        employee.role = next_role
        employee.salary_dollars += next_role.value
        print(f"{employee.name} was promoted to a package of {employee.salary_dollars}.")

However, you could ensure that a new Role doesn't need to be put in order by using sorted:

class Company:
    def promote_employee(self, employee: Employee):
        # sort by the values
        roles = sorted(Role, key=lambda x: x.value)

        # now use the sorted 'roles'
        next_role = next((role for role in roles if role.value > employee.role.value), None)
        ...

As an aside, for an even more concise sorted call, you could use the operator.attrgetter function:

from operator import attrgetter

sorted(Role, key=attrgetter('value'))
[<Role.INTERN: 0>, <Role.WORKER: 5000>, <Role.LEAD: 15000>, <Role.MANAGER: 25000>]

One last tidbit: the sorted Role doesn't really change from call to call. It might be better to make this a classvariable. To do this in a dataclass, you can use a type signature:

from typing import ClassVar


@dataclass
class Company:
    ~snip~
    roles: ClassVar[list] = sorted(Role, key=attrgetter('value')) # this gets set once
    
    def promote_employee(self, employee: Employee):
        # now use the sorted 'roles'
        next_role = next((role for role in self.roles if role.value > employee.role.value), None)
    

accept_resignation

This should just be a call to Company.remove_employee:

class Employee:
    ~snip~
    def resign(self) -> Person:
        self.company.remove_employee(self)
        person_attrs = dict(**self)
        
        for k in ('company', 'id', 'role', 'team', 'working_days', 'leave_days'):
            person_attrs.pop(k)
            
        return Person(**person_attrs)

accept_resignation

This should just be a call to Company.remove_employee:

class Employee:
    ~snip~
    def resign(self) -> Person:
        self.company.remove_employee(self)
        person_attrs = dict(**self)
        
        for k in EXTRA_FIELDS:
            person_attrs.pop(k)
            
        return Person(**person_attrs)

But if we look at the implementation of Company.fire_employee, it is identical, plus it gives the convenience of returning a Person instance.

@dataclass
class Company():
    ~snip~
    def fire_employee(self, employee: Employee) -> Person:
        if employee.name not in self.employees:
            raise ValueError(f"Employee {employee.name} not found in company")

        self.remove_employee(employee)

        attrs = dict(employee)

        for k in EXTRA_FIELDS:
            attrs.pop(k)

        return Person(**attrs)

So let's call that instead:

class Employee:
    ~snip~
    def resign(self) -> Person:
        return self.company.fire_employee(self)

It looks weird, because the employee isn't technically being fired, but it certainly makes the interface easy to reason about.

Final Result

from dataclasses import fields, dataclass
from enum import auto, Enum
from operator import attrgetter
from typing import ClassVar


class Role(Enum):
    INTERN = 0
    WORKER = 5000
    LEAD = 15000
    MANAGER = 25000
    
    
class Team(Enum):
    MARKETING = auto()
    PRODUCT = auto()
    DESIGN = auto()
    SALES = auto()
    FINANCE = auto()
    
    
@dataclass
class Company:
    name: str
    minimum_hiring_grades: int
    minimum_work_experience: int
    employee_leaves: int
    employee_bonus_percent: int
    employee_working_days: int
    employees: dict = field(init=False, default_factory=dict)
    teams: dict = field(init=False)
    roles: ClassVar[list] = sorted(Role, key=attrgetter('value'))
    counter: int = 1
        
    def __post_init__(self):
        # also iterate directly over your Enums
        # here
        self.teams = {
            team.name: {
                role.name: None 
                for role in Role
            } for team in Team
        }
        
    def create_id(self) -> str:
        """Creates an id like ac00001 for the first employee at AcmeCo"""
        abbrev = self.name[:2].lower()
        num = str(self.counter).zfill(5)
        self.counter += 1
        return ''.join((abbrev, num))
    
        
    def add_employee(self, employee: 'Employee') -> None:
        self.employees[employee.name] = employee
        self.teams[employee.team.name][employee.role.name] = employee
        
    
    def remove_employee(self, employee: 'Employee') -> None:
        self.employees.pop(employee.name)
        self.teams[employee.team.name][employee.role.name] = None
        
        
    def role_occupied(self, role: Role, team: Team) -> bool:
        return bool(self.teams.get(team.name, {}).get(role.name))
    
    
    def is_applicant_eligible(self, applicant: Person, role: Role, team: Team) -> bool:
        employee_has_work_experience = applicant.work_experience_years >= self.minimum_work_experience
        employee_has_good_grades = applicant.grades_percent_average >= self.minimum_hiring_grades
        
        if self.role_occupied(role, team):
            return False
        elif not employee_has_work_experience and not employee_has_good_grades:
            return False
        else:
            return True
    
    
    def hire_employee(self, person: 'Person', role: Role, team: Team) -> 'Employee':
        salary = (person.grades_percent_average + person.work_experience_years) * 1000
        
        employee = Employee(
            company=self,
            role=role,
            team=team,
            salary=salary,
            id=self.create_id(),
            **person
        )
        
        self.add_employee(employee)
        return employee
    
     
    def promote_employee(self, employee: 'Employee') -> None:
        next_role = next(
            (role for role in self.roles if role.value > employee.role.value), 
            None
        )

        if next_role is None:
            raise Exception("Can't promote a manager!")

        employee.role = next_role
        employee.salary_dollars += next_role.value
        print(f"{employee.name} was promoted to a package of {employee.salary_dollars}.")

    
    def fire_employee(self, employee: 'Employee') -> 'Person':
        if employee.name not in self.employees:
            raise ValueError(f"Employee {employee.name} not found in company")

        self.remove_employee(employee)

        attrs = dict(employee)

        for k in EXTRA_FIELDS:
            attrs.pop(k)

        return Person(**attrs)
    
    

@dataclass
class Person:
    name : str
    grades_percent_average : int
    work_experience_years : int
        
    def keys(self):
        return [field.name for field in fields(self)]
    
    def __getitem__(self, key):
        return getattr(self, key)
    
    def apply_for_role(self, company: Company, role: Role, team: Team) -> 'Person':
        if not company.is_applicant_eligible(self, role=role, team=team):
            return self
        elif company.role_occupied(role=role, team=team):
            return self
        else:
            return company.hire_employee(self, role, team)
    

@dataclass
class Employee(Person):
    company: Company
    working_days: int = field(init=False)
    leave_days: int = field(init=False)
    role: Role
    team: Team
    salary: int
    id: str
        
    def __post_init__(self):
        self.working_days = self.company.employee_working_days
        self.leave_days = self.company.employee_leaves
        
        
    def request_leave(self, leaves_required: int) -> None:
        too_few_required = leaves_required <= 3 
        too_few_remaining = (employee.available_leaves - leaves_required) <= 0
        
        if too_few_required and too_few_remaining:
            print("Not enough leave available")
            return
        
        self.working_days -= leaves_required
        self.leave_days -= leaves_required
        
        
    def resign(self) -> Person:
        return self.company.fire_employee(self)
    
    
    
EXTRA_FIELDS = list({f.name for f in fields(Employee)} - {f.name for f in fields(Person)})

And to show basic usage

b = Company(name="b", minimum_hiring_grades=90, minimum_work_experience=5, employee_leaves=30, employee_bonus_percent=25, employee_working_days=300)
arham = Person(name="Arham", grades_percent_average=92, work_experience_years=5)
rajeev = Person(name="Rajeev", grades_percent_average=91, work_experience_years=5)

print(arham)
Person(name='Arham', grades_percent_average=92, work_experience_years=5)

arham = arham.apply_for_role(b, Role.INTERN, Team.MARKETING)

print(arham)
Employee(name='Arham', grades_percent_average=92, work_experience_years=5, company=Company(name='b', minimum_hiring_grades=90, minimum_work_experience=5, employee_leaves=30, employee_bonus_percent=25, employee_working_days=300, employees={'Arham': ...}, teams={'MARKETING': {'INTERN': ..., 'WORKER': None, 'LEAD': None, 'MANAGER': None}, 'PRODUCT': {'INTERN': None, 'WORKER': None, 'LEAD': None, 'MANAGER': None}, 'DESIGN': {'INTERN': None, 'WORKER': None, 'LEAD': None, 'MANAGER': None}, 'SALES': {'INTERN': None, 'WORKER': None, 'LEAD': None, 'MANAGER': None}, 'FINANCE': {'INTERN': None, 'WORKER': None, 'LEAD': None, 'MANAGER': None}}, counter=1), working_days=300, leave_days=30, role=<Role.INTERN: 0>, team=<Team.MARKETING: 1>, salary=97000, id='b00001')

New Bug!

If arham applies for the same job twice, it makes it appear that he's now unemployed:

arham = arham.apply_for_role(b, Role.INTERN, Team.MARKETING)
print(arham)
Person(name='Arham', grades_percent_average=92, work_experience_years=5)

So let's make sure when being accepted for a new role, they are released from the old one:

@dataclass
class Employee(Person):   
    def apply_for_role(self, company: Company, role: Role, team: Team) -> 'Employee':
        if not company.is_applicant_eligible(self, role=role, team=team):
            return self
        elif company.role_occupied(role=role, team=team):
            return self
        
        person = self.resign()
        
        return company.hire_employee(person, role=role, team=team)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.