Coverage for o2/models/timetable/resource_pool.py: 100%

1from dataclasses import dataclass, replace 

2 

3from dataclass_wizard import JSONWizard 

4 

5from o2.models.timetable.resource import Resource 

6 

7 

8@dataclass(frozen=True) 

9class ResourcePool(JSONWizard): 

10 """Collection of resources that can be assigned to a task.""" 

11 

12 id: str 

13 name: str 

14 resource_list: list[Resource] 

15 fixed_cost_fn: str = "0" 

16 

17 def remove_resource(self, resource_id: str) -> "ResourcePool": 

18 """Remove a resource from the pool.""" 

19 return replace( 

20 self, 

21 resource_list=[resource for resource in self.resource_list if resource.id != resource_id], 

22 ) 

23 

24 def update_resource(self, updated_resource: Resource) -> "ResourcePool": 

25 """Update a resource in the pool.""" 

26 return replace( 

27 self, 

28 resource_list=[ 

29 updated_resource if updated_resource.id == resource.id else resource 

30 for resource in self.resource_list 

31 ], 

32 )