Coverage for backend/django/core/auxiliary/viewsets/CompoundViewSet.py: 100%

1from drf_spectacular.utils import extend_schema 

2from pydantic import BaseModel 

3from rest_framework import status, viewsets 

4from rest_framework.response import Response 

5from ahuora_compounds.CompoundDB import db 

6from ahuora_compounds.packages.chemsep import ChemsepCompoundData 

7 

8class CompoundListItem(BaseModel): 

9 name: str 

10 formula: str | None = None 

11 

12class CompoundViewSet(viewsets.ViewSet): 

13 """Viewset for retrieving compounds and their physical properties.""" 

14 

15 @extend_schema(responses={status.HTTP_200_OK: CompoundListItem}) 

16 def list(self, request): 

17 """List all compounds.""" 

18 

19 def get_compound_item(compound_name: str) -> CompoundListItem: 

20 compound_info = ChemsepCompoundData.get(compound_name, None) 

21 

22 return CompoundListItem( 

23 name=compound_name, 

24 formula=compound_info.StructureFormula.value if compound_info else None, 

25 ) 

26 

27 compounds = [ 

28 get_compound_item(compound_name).model_dump() 

29 for compound_name in 

30 db.get_compound_names() 

31 ] 

32 

33 return Response(compounds, status=status.HTTP_200_OK)