Coverage for backend/django/core/auxiliary/viewsets/DataColumnViewSet.py: 96%

1from core.viewset import ModelViewSet 

2from rest_framework.response import Response 

3from core.auxiliary.models.DataColumn import DataColumn 

4from core.auxiliary.serializers.DataColumnSerializer import DataColumnSerializer 

5from drf_spectacular.utils import extend_schema, OpenApiParameter 

6from drf_spectacular.types import OpenApiTypes 

7from rest_framework.decorators import action 

8from rest_framework.serializers import DictField, IntegerField 

9from django.db.models import QuerySet 

10class DataColumnViewSet(ModelViewSet): 

11 serializer_class = DataColumnSerializer 

12 

13 def get_queryset(self): 

14 queryset = DataColumn.objects.select_related( 

15 "property_value__property__set__simulationObject", 

16 ) 

17 return queryset 

18 

19 @extend_schema( 

20 parameters=[ 

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

22 ] 

23 ) 

24 def list(self, request): 

25 queryset = self.get_queryset() 

26 return Response(DataColumnSerializer(queryset.filter(scenario_id=request.query_params.get("scenario")), many=True).data) 

27 

28 

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

30 parameters=[ 

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

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

33 ]) 

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

35 def download_tag_mappings(self, request): 

36 return Response( 

37 create_data_column_json(self.get_queryset()), 

38 content_type="text/json", 

39 headers={'Content-Disposition': 'attachment; filename="mss-tag-mappings.json"'} 

40 ) 

41 

42def create_data_column_json(queryset: QuerySet[DataColumn, DataColumn]) -> dict[str, dict[str, str]]: 

43 """ 

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

45 The keys are the column names, and the values are the IDs of the associated PropertyValues. 

46 """ 

47 data_column: DataColumn 

48 data = { 

49 data_column.name: { 

50 "property": data_column.property_value.pk, 

51 "units": data_column.property_value.property.unit 

52 } for data_column in queryset.all() 

53 } 

54 return data