From 2fbffb4b3bd73d7e496886a439b8a8f9fed54bec Mon Sep 17 00:00:00 2001
From: Tim Lorsbach <tim@lorsba.ch>
Date: Fri, 3 Jul 2026 09:40:47 +0200
Subject: [PATCH] Fixed molfile assignment, adjusted Export

---
 bayer/models.py   |  4 ++--
 utilities/misc.py | 42 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/bayer/models.py b/bayer/models.py
index 0ee478c4..afa61f2e 100644
--- a/bayer/models.py
+++ b/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
 
diff --git a/utilities/misc.py b/utilities/misc.py
index e1a84dec..7208c0df 100644
--- a/utilities/misc.py
+++ b/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
-    compounds: List[CompoundExportSchema]
+    license: LicenseExportSchema | None = None
+    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):