Coverage for backend/core/auxiliary/methods/copy_flowsheet/copy_caching.py: 100%

1from django.db.models import Model 

2from dataclasses import dataclass 

3from typing import List, Dict, Type 

4from core.validation import flowsheet_ctx 

5 

6 

7class ModelLookup: 

8 """ 

9 A class to help map model primary keys to their respective model classes. 

10 """ 

11 

12 def __init__(self, object_list: List[Model]): 

13 self.model_map: Dict[int, Model] = {} 

14 for model in object_list: 

15 self.model_map[model.pk] = model 

16 

17 def get_model(self, pk: int) -> Model: 

18 return self.model_map.get(pk, None) 

19 

20 # Enable iteration over the model_map values 

21 def __iter__(self): 

22 return iter(self.model_map.values()) 

23 

24 

25@dataclass 

26class RelatedModels: 

27 model: Model 

28 related_model_ids: List[int] 

29 

30 

31def get_prefetch_fields(model: Type[Model]) -> List[str]: 

32 """ 

33 We want to prefetch all fields that are many to many relationships 

34 

35 """ 

36 prefetch_fields = [] 

37 for field in model._meta.get_fields(): 

38 if field.is_relation and field.many_to_many and not field.auto_created: 

39 prefetch_fields.append(field.name) 

40 return prefetch_fields 

41 

42 

43ModelLookupDict = dict[Type[Model], ModelLookup]