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

1from drf_spectacular.utils import extend_schema 

2from pydantic import BaseModel 

3from rest_framework import status, viewsets 

4from rest_framework.response import Response 

5from compounds.CompoundDB import db 

6 

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 = db.get_compound(compound_name) 

21 compound_info = compound.get_source("chemsep") 

22 

23 return CompoundListItem( 

24 name=compound_name, 

25 formula=compound_info.get("StructureFormula", None).value if compound_info else None, 

26 ) 

27 

28 compounds = [ 

29 get_compound_item(compound_name).model_dump() 

30 for compound_name in 

31 db.get_compound_names() 

32 ] 

33 

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