Fixed molfile assignment, adjusted Export

bayer/models.py

@@ -185,7 +185,7 @@ class PESStructure(CompoundStructure):
     def create(
             compound: Compound,
             pes_link: str,
-            mol_file: str,
+            molfile: str,
             smiles: str,
             name: str = None,
             description: str = None,
@@ -204,7 +204,7 @@ class PESStructure(CompoundStructure):
             cs.description = nh3.clean(description, tags=s.ALLOWED_HTML_TAGS).strip()
 
         cs.smiles = smiles
-        cs.mol_file = mol_file
+        cs.molfile = molfile
         cs.pes_link = pes_link
         cs.compound = compound
 

utilities/misc.py

@@ -13,6 +13,7 @@ from django.db import transaction
 from ninja import Schema
 from pydantic import HttpUrl
 
+from bayer.models import PESCompound
 from epdb.models import (
     AdditionalInformation,
     Compound,
@@ -64,6 +65,9 @@ class RefExportSchema(Schema):
         return str(value)
 
 
+class RefGroupExportSchema(RefExportSchema): ...
+
+
 class RefCompoundExportSchema(RefExportSchema): ...
 
 
@@ -110,6 +114,14 @@ class CompoundStructureExportSchema(RefCompoundStructureExportSchema):
     scenarios: List[RefScenarioExportSchema]
 
 
+class PESCompoundExportSchema(CompoundExportSchema):
+    structures: List["PESCompoundStructureExportSchema"]
+
+
+class PESCompoundStructureExportSchema(CompoundStructureExportSchema):
+    pes_link: HttpUrl
+
+
 ############
 # Reaction #
 ############
@@ -210,8 +222,10 @@ class PackageExportSchema(Schema):
     description: str
     uuid: str
     reviewed: bool
-    license: LicenseExportSchema | None
+    license: LicenseExportSchema | None = None
-    compounds: List[CompoundExportSchema]
+    classification_level: str = "Internal"
+    data_pool: RefGroupExportSchema | None = None
+    compounds: List[CompoundExportSchema | PESCompoundExportSchema]
     reactions: List[ReactionExportSchema]
     simple_rules: List[RuleExportSchema]
     composite_rules: List[ParallelRuleExportSchema]
@@ -224,6 +238,30 @@ class PackageExportSchema(Schema):
         value = obj.get("uuid") if isinstance(obj, dict) else obj.uuid
         return str(value)
 
+    @staticmethod
+    def resolve_classification_level(obj):
+        if isinstance(obj, str):
+            return obj["classification_level"]
+        return obj.Classification(obj.classification_level).name
+
+
+    @staticmethod
+    def resolve_compounds(obj):
+        res = []
+        if isinstance(obj, dict):
+            for c in obj.get("compounds", []):
+                if "pes_link" in c:
+                    res.append(PESCompoundExportSchema.model_validate(c))
+                else:
+                    res.append(CompoundExportSchema.model_validate(c))
+        else:
+            for c in obj.compounds.all():
+                if isinstance(c, PESCompound):
+                    res.append(PESCompoundExportSchema.from_orm(c))
+                else:
+                    res.append(CompoundExportSchema.from_orm(c))
+        return res
+
     @staticmethod
     def resolve_simple_rules(obj):
         if isinstance(obj, dict):