Coverage for backend/django/core/auxiliary/viewsets/ExpressionViewSet.py: 68%

1from core.viewset import ModelViewSet 

2from django.db.models import QuerySet 

3from rest_framework.response import Response 

4from core.auxiliary.models.Expression import Expression 

5from core.auxiliary.serializers.ExpressionSerializer import ExpressionSerializer 

6from drf_spectacular.utils import extend_schema, OpenApiParameter, OpenApiTypes 

7from rest_framework.decorators import action 

8from rest_framework.serializers import DictField, IntegerField 

9 

10 

11class ExpressionViewSet(ModelViewSet): 

12 serializer_class = ExpressionSerializer 

13 

14 # Filter to only show expressions for a specific flowsheet 

15 def get_queryset(self) -> QuerySet[Expression]: 

16 queryset = Expression.objects.all() 

17 scenario_id = self.request.query_params.get("scenario") 

18 if scenario_id is not None: 

19 queryset = queryset.filter(scenario_id=scenario_id) 

20 return queryset 

21 

22 @extend_schema( 

23 parameters=[ 

24 OpenApiParameter(name="scenario", required=True, type=OpenApiTypes.INT), 

25 ] 

26 ) 

27 def list(self, request): 

28 return super().list(request) 

29 

30 

31 @extend_schema(responses=DictField(child=IntegerField()), 

32 parameters=[ 

33 OpenApiParameter(name="scenario", required=True, type=OpenApiTypes.INT), 

34 OpenApiParameter(name="flowsheet", required=True, type=OpenApiTypes.INT), 

35 ]) 

36 @action(detail=False,methods=["get"]) 

37 def download_tag_mappings(self, request): 

38 return Response(create_expression_json(self.get_queryset()), 

39 content_type="text/json", 

40 headers={ 

41 'Content-Disposition': 'attachment; filename="mss-tag-mappings.json"' 

42 }) 

43 

44 

45def create_expression_json(queryset: QuerySet[Expression, Expression]) -> dict[str, int]: 

46 """ 

47 Create a JSON representation of the MSS columns in the scenario. 

48 The keys are the expression names, and the values are the IDs of the associated PropertyValues. 

49 """ 

50 expr: Expression 

51 data = { 

52 expr.name: {"property": expr.property.values.first().pk, 

53 "units": expr.property.unit 

54 } for expr in queryset.all() 

55 } 

56 return data