Fix Classification import, Fix PES Deserialization

utilities/misc.py

@@ -36,7 +36,7 @@ from epdb.models import (
     Setting,
     SimpleAmbitRule,
     SimpleRDKitRule,
-    SimpleRule,
+    SimpleRule, Group,
 )
 from utilities.chem import FormatConverter
 
@@ -44,7 +44,7 @@ logger = logging.getLogger(__name__)
 Package = s.GET_PACKAGE_MODEL()
 
 if TYPE_CHECKING:
-    from epdb.logic import SPathway
+    from epdb.logic import SPathway, GroupManager
 
 
 class LicenseExportSchema(Schema):
@@ -123,8 +123,8 @@ class PESCompoundExportSchema(RefCompoundExportSchema):
     description: str
     aliases: List[str]
     default_structure: RefCompoundStructureExportSchema
-    structures: List["CompoundStructureExportSchema"]
     structures: List["PESCompoundStructureExportSchema"]
+    scenarios: List[RefScenarioExportSchema]
 
 
 class PESCompoundStructureExportSchema(RefCompoundStructureExportSchema):
@@ -277,7 +277,7 @@ class PackageExportSchema(Schema):
 
     @staticmethod
     def resolve_classification_level(obj):
-        if isinstance(obj, str):
+        if isinstance(obj, dict):
             return obj["classification_level"]
         return obj.Classification(obj.classification_level).name
 
@@ -602,15 +602,20 @@ class PackageImporter:
         package_license = License.objects.get(link=parsed.license.link) if parsed.license else None
         package_classification_level = Package.Classification[parsed.classification_level.upper()]
 
-        package = Package.objects.create(
-            uuid=package_uuid,
-            name=f"{parsed.name} - Imported at {datetime.now()}",
-            description=parsed.description,
-            reviewed=False, # Everything will be imported as non reviewed
-            license=package_license,
-            classification_level=package_classification_level,
-            data_pool=parsed.data_pool,
-        )
+        print(parsed.classification_level, "->", package_classification_level)
+
+        package = Package()
+        package.uuid = package_uuid
+        package.name = f"{parsed.name} - Imported at {datetime.now()}"
+        package.description = parsed.description
+        package.reviewed = False
+        package.license = package_license
+        package.classification_level = package_classification_level
+
+        if package_classification_level == Package.Classification.SECRET:
+            package.data_pool = Group.objects.get(uuid=parsed.data_pool.uuid)
+
+        package.save()
 
         # Import Elements
         self._import_compounds(package, parsed.compounds)