forked from enviPath/enviPy
Package Export/Import cycle
This commit is contained in:
Tim Lorsbach
@ -1,2 +1,5 @@
|
||||
class InvalidSMILESException(Exception):
|
||||
pass
|
||||
|
||||
class PackageImportException(Exception):
|
||||
pass
|
||||
@ -1078,10 +1078,8 @@ class PackageManager(object):
|
||||
data: Dict[str, Any],
|
||||
owner: User,
|
||||
preserve_uuids=False,
|
||||
add_import_timestamp=True,
|
||||
trust_reviewed=False,
|
||||
) -> Package:
|
||||
importer = PackageImporter(data, preserve_uuids, add_import_timestamp, trust_reviewed)
|
||||
importer = PackageImporter(data, preserve_uuids)
|
||||
imported_package = importer.do_import()
|
||||
|
||||
up = UserPackagePermission()
|
||||
|
||||
@ -11,9 +11,10 @@ from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
from django.conf import settings as s
|
||||
from django.db import transaction
|
||||
from ninja import Schema
|
||||
from pydantic import HttpUrl
|
||||
from pydantic import HttpUrl, ValidationError
|
||||
|
||||
from bayer.models import PESCompound
|
||||
from bayer.models import PESCompound, PESStructure
|
||||
from epdb.exceptions import PackageImportException
|
||||
from epdb.models import (
|
||||
AdditionalInformation,
|
||||
Compound,
|
||||
@ -92,6 +93,9 @@ class RefPathwayExportSchema(RefExportSchema): ...
|
||||
class RefScenarioExportSchema(RefExportSchema): ...
|
||||
|
||||
|
||||
class RefEnzymeExportSchema(RefExportSchema): ...
|
||||
|
||||
|
||||
############
|
||||
# Compound #
|
||||
############
|
||||
@ -114,11 +118,23 @@ class CompoundStructureExportSchema(RefCompoundStructureExportSchema):
|
||||
scenarios: List[RefScenarioExportSchema]
|
||||
|
||||
|
||||
class PESCompoundExportSchema(CompoundExportSchema):
|
||||
class PESCompoundExportSchema(RefCompoundExportSchema):
|
||||
name: str
|
||||
description: str
|
||||
aliases: List[str]
|
||||
default_structure: RefCompoundStructureExportSchema
|
||||
structures: List["CompoundStructureExportSchema"]
|
||||
structures: List["PESCompoundStructureExportSchema"]
|
||||
|
||||
|
||||
class PESCompoundStructureExportSchema(CompoundStructureExportSchema):
|
||||
class PESCompoundStructureExportSchema(RefCompoundStructureExportSchema):
|
||||
name: str
|
||||
description: str
|
||||
aliases: List[str]
|
||||
smiles: str
|
||||
molfile: Optional[str]
|
||||
normalized_structure: bool
|
||||
scenarios: List[RefScenarioExportSchema]
|
||||
pes_link: HttpUrl
|
||||
|
||||
|
||||
@ -140,7 +156,28 @@ class ReactionExportSchema(RefReactionExportSchema):
|
||||
#########
|
||||
# Rules #
|
||||
#########
|
||||
class RuleExportSchema(RefRuleExportSchema):
|
||||
class EnzymeExportSchema(RefEnzymeExportSchema):
|
||||
ec_number: str
|
||||
classification_level: int
|
||||
linking_method: str
|
||||
reaction_evidence: List[RefReactionExportSchema]
|
||||
edge_evidence: List[RefEdgeExportSchema]
|
||||
|
||||
|
||||
class EnzymeRuleExportSchema(RefRuleExportSchema):
|
||||
enzymes: List[EnzymeExportSchema] | None = None
|
||||
|
||||
@staticmethod
|
||||
def resolve_enzymes(obj):
|
||||
if isinstance(obj, dict):
|
||||
res = []
|
||||
for e in obj.get("enzymes", []):
|
||||
res.append(EnzymeExportSchema.model_validate(e))
|
||||
return res
|
||||
return obj.enzymelink_set.all()
|
||||
|
||||
|
||||
class RuleExportSchema(EnzymeRuleExportSchema):
|
||||
name: str
|
||||
description: str
|
||||
aliases: List[str]
|
||||
@ -150,7 +187,7 @@ class RuleExportSchema(RefRuleExportSchema):
|
||||
scenarios: List[RefScenarioExportSchema]
|
||||
|
||||
|
||||
class ParallelRuleExportSchema(RefRuleExportSchema):
|
||||
class ParallelRuleExportSchema(EnzymeRuleExportSchema):
|
||||
name: str
|
||||
description: str
|
||||
aliases: List[str]
|
||||
@ -250,7 +287,8 @@ class PackageExportSchema(Schema):
|
||||
res = []
|
||||
if isinstance(obj, dict):
|
||||
for c in obj.get("compounds", []):
|
||||
if "pes_link" in c:
|
||||
is_pes = any([cs.get("pes_link", None) is not None for cs in c.get("structures", [])])
|
||||
if is_pes:
|
||||
res.append(PESCompoundExportSchema.model_validate(c))
|
||||
else:
|
||||
res.append(CompoundExportSchema.model_validate(c))
|
||||
@ -315,51 +353,276 @@ class PackageExporter:
|
||||
|
||||
|
||||
class PackageImporter:
|
||||
"""
|
||||
Imports package data from JSON export.
|
||||
Handles object creation, relationship mapping, and dependency resolution.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
package: Dict[str, Any],
|
||||
preserve_uuids: bool = False,
|
||||
add_import_timestamp=True,
|
||||
trust_reviewed=False,
|
||||
):
|
||||
"""
|
||||
Initialize the importer.
|
||||
|
||||
Args:
|
||||
preserve_uuids: If True, preserve original UUIDs. If False, generate new ones.
|
||||
"""
|
||||
def __init__(self, package: Dict[str, Any], preserve_uuids: bool = False):
|
||||
self.preserve_uuids = preserve_uuids
|
||||
self.add_import_timestamp = add_import_timestamp
|
||||
self.trust_reviewed = trust_reviewed
|
||||
self.uuid_mapping = {}
|
||||
self.object_cache = {}
|
||||
self._raw_package = package
|
||||
|
||||
def _get_or_generate_uuid(self, original_uuid: str) -> str:
|
||||
"""Get mapped UUID or generate new one if not preserving UUIDs."""
|
||||
if self.preserve_uuids:
|
||||
return original_uuid
|
||||
|
||||
if original_uuid not in self.uuid_mapping:
|
||||
self.uuid_mapping[original_uuid] = str(uuid.uuid4())
|
||||
|
||||
return self.uuid_mapping[original_uuid]
|
||||
|
||||
def _cache_object(self, model_name: str, uuid_str: str, obj):
|
||||
"""Cache a created object for later reference."""
|
||||
self.object_cache[(model_name, uuid_str)] = obj
|
||||
|
||||
def _get_cached_object(self, model_name: str, uuid_str: str):
|
||||
"""Get a cached object by model name and UUID."""
|
||||
return self.object_cache.get((model_name, uuid_str))
|
||||
self._cache = {}
|
||||
|
||||
def do_import(self) -> Package:
|
||||
return self._import_package_from_json(self._raw_package)
|
||||
return self._import_package_from_json()
|
||||
|
||||
def _import_compounds(
|
||||
self,
|
||||
package: Package,
|
||||
compounds: List[CompoundExportSchema | PESCompoundExportSchema],
|
||||
):
|
||||
for c in compounds:
|
||||
|
||||
if isinstance(c, PESCompoundExportSchema):
|
||||
c_cls = PESCompound
|
||||
else:
|
||||
c_cls = Compound
|
||||
|
||||
new_c = c_cls()
|
||||
new_c.uuid = str(uuid.uuid4()) if not self.preserve_uuids else c.uuid
|
||||
new_c.package = package
|
||||
new_c.name = c.name
|
||||
new_c.description = c.description
|
||||
new_c.aliases = c.aliases
|
||||
new_c.save()
|
||||
|
||||
self._cache[c.uuid] = new_c
|
||||
|
||||
for cs in c.structures:
|
||||
|
||||
if isinstance(cs, PESCompoundStructureExportSchema):
|
||||
cs_cls = PESStructure
|
||||
else:
|
||||
cs_cls = CompoundStructure
|
||||
|
||||
new_cs = cs_cls()
|
||||
new_cs.uuid = str(uuid.uuid4()) if not self.preserve_uuids else cs.uuid
|
||||
new_cs.compound = new_c
|
||||
new_cs.name = cs.name
|
||||
new_cs.description = cs.description
|
||||
new_cs.aliases = cs.aliases
|
||||
new_cs.smiles = cs.smiles
|
||||
new_cs.molfile = cs.molfile
|
||||
|
||||
if cs_cls == PESStructure:
|
||||
new_cs.pes_link = cs.pes_link
|
||||
|
||||
new_cs.normalized_structure = cs.normalized_structure
|
||||
new_cs.save()
|
||||
|
||||
self._cache[cs.uuid] = new_cs
|
||||
|
||||
# Now we can assigne the default_structure to the compound
|
||||
new_c.default_structure = self._cache[c.default_structure.uuid]
|
||||
new_c.save()
|
||||
|
||||
def _import_simple_rules(
|
||||
self,
|
||||
package: Package,
|
||||
rules: List[RuleExportSchema]
|
||||
):
|
||||
for r in rules:
|
||||
new_r = SimpleAmbitRule()
|
||||
new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid
|
||||
new_r.package = package
|
||||
new_r.name = r.name
|
||||
new_r.description = r.description
|
||||
new_r.aliases = r.aliases
|
||||
new_r.smirks = r.smirks
|
||||
new_r.reactant_filter_smarts = r.reactant_filter_smarts
|
||||
new_r.product_filter_smarts = r.product_filter_smarts
|
||||
new_r.save()
|
||||
|
||||
self._cache[r.uuid] = new_r
|
||||
|
||||
def _import_composite_rules(
|
||||
self,
|
||||
package: Package,
|
||||
rules: List[ParallelRuleExportSchema]
|
||||
):
|
||||
for r in rules:
|
||||
new_r = ParallelRule()
|
||||
new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid
|
||||
new_r.package = package
|
||||
new_r.name = r.name
|
||||
new_r.description = r.description
|
||||
new_r.aliases = r.aliases
|
||||
new_r.save()
|
||||
|
||||
for sr in r.simple_rules:
|
||||
new_r.simple_rules.add(self._cache[sr.uuid])
|
||||
|
||||
self._cache[r.uuid] = new_r
|
||||
|
||||
def _import_reactions(
|
||||
self,
|
||||
package: Package,
|
||||
reactions: List[ReactionExportSchema],
|
||||
):
|
||||
for r in reactions:
|
||||
new_r = Reaction()
|
||||
new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid
|
||||
new_r.package = package
|
||||
new_r.name = r.name
|
||||
new_r.description = r.description
|
||||
new_r.aliases = r.aliases
|
||||
new_r.multi_step = r.multi_step
|
||||
new_r.medline_references = r.medline_references
|
||||
new_r.save()
|
||||
|
||||
for educt in r.educts:
|
||||
new_r.educts.add(self._cache[educt.uuid])
|
||||
|
||||
for product in r.products:
|
||||
new_r.products.add(self._cache[product.uuid])
|
||||
|
||||
for rule in r.rules:
|
||||
new_r.rules.add(self._cache[rule.uuid])
|
||||
|
||||
self._cache[r.uuid] = new_r
|
||||
|
||||
def _import_pathways(
|
||||
self,
|
||||
package: Package,
|
||||
pathways: List[PathwayExportSchema]
|
||||
):
|
||||
for pw in pathways:
|
||||
new_pw = Pathway()
|
||||
new_pw.uuid = str(uuid.uuid4()) if not self.preserve_uuids else pw.uuid
|
||||
new_pw.package = package
|
||||
new_pw.name = pw.name
|
||||
new_pw.description = pw.description
|
||||
new_pw.aliases = pw.aliases
|
||||
new_pw.predicted = pw.predicted
|
||||
new_pw.save()
|
||||
|
||||
self._cache[pw.uuid] = new_pw
|
||||
|
||||
for n in pw.nodes:
|
||||
new_n = Node()
|
||||
new_n.uuid = str(uuid.uuid4()) if not self.preserve_uuids else n.uuid
|
||||
new_n.pathway = new_pw
|
||||
new_n.name = n.name
|
||||
new_n.description = n.description
|
||||
new_n.aliases = n.aliases
|
||||
new_n.default_node_label = self._cache[n.default_node_label.uuid]
|
||||
new_n.depth = n.depth
|
||||
new_n.stereo_removed = n.stereo_removed
|
||||
new_n.save()
|
||||
|
||||
for nl in n.node_labels:
|
||||
new_n.node_labels.add(self._cache[nl.uuid])
|
||||
|
||||
self._cache[n.uuid] = new_n
|
||||
|
||||
for e in pw.edges:
|
||||
new_e = Edge()
|
||||
new_e.uuid = str(uuid.uuid4()) if not self.preserve_uuids else e.uuid
|
||||
new_e.pathway = new_pw
|
||||
new_e.name = e.name
|
||||
new_e.description = e.description
|
||||
new_e.aliases = e.aliases
|
||||
new_e.edge_label = self._cache[e.edge_label.uuid]
|
||||
new_e.save()
|
||||
|
||||
for sn in e.start_nodes:
|
||||
new_e.start_nodes.add(self._cache[sn.uuid])
|
||||
|
||||
for en in e.end_nodes:
|
||||
new_e.end_nodes.add(self._cache[en.uuid])
|
||||
|
||||
self._cache[e.uuid] = new_e
|
||||
|
||||
def _import_scenarios(
|
||||
self,
|
||||
package: Package,
|
||||
scenarios: List[ScenarioExportSchema]
|
||||
):
|
||||
for s in scenarios:
|
||||
new_s = Scenario()
|
||||
new_s.uuid = str(uuid.uuid4()) if not self.preserve_uuids else s.uuid
|
||||
new_s.package = package
|
||||
new_s.name = s.name
|
||||
new_s.description = s.description
|
||||
new_s.scenario_date = s.scenario_date
|
||||
new_s.scenario_type = s.scenario_type
|
||||
new_s.save()
|
||||
|
||||
self._cache[s.uuid] = new_s
|
||||
|
||||
def _import_additional_information(
|
||||
self,
|
||||
package: Package,
|
||||
additional_information: List[AdditionalInformationExportSchema]
|
||||
):
|
||||
for ai in additional_information:
|
||||
new_ai = AdditionalInformation()
|
||||
new_ai.uuid = str(uuid.uuid4()) if not self.preserve_uuids else ai.uuid
|
||||
new_ai.package = package
|
||||
new_ai.type = ai.type
|
||||
new_ai.data = ai.data
|
||||
|
||||
if ai.scenario:
|
||||
new_ai.scenario = self._cache[ai.scenario.uuid]
|
||||
|
||||
if ai.attach_object:
|
||||
new_ai.content_object = self._cache[ai.attach_object.uuid]
|
||||
|
||||
new_ai.save()
|
||||
|
||||
self._cache[ai.uuid] = new_ai
|
||||
|
||||
def _link_scenarios_after_import(self, data: PackageExportSchema):
|
||||
collections = [
|
||||
data.compounds,
|
||||
data.simple_rules,
|
||||
data.composite_rules,
|
||||
data.reactions,
|
||||
data.pathways,
|
||||
[n for pw in data.pathways for n in pw.nodes],
|
||||
[e for pw in data.pathways for e in pw.edges],
|
||||
]
|
||||
|
||||
for coll in collections:
|
||||
for elem in coll:
|
||||
elem_obj = self._cache[elem.uuid]
|
||||
for s in elem.scenarios:
|
||||
elem_obj.scenarios.add(self._cache[s.uuid])
|
||||
|
||||
|
||||
def _import_package_from_json(
|
||||
self,
|
||||
) -> Package:
|
||||
try:
|
||||
parsed = PackageExportSchema.model_validate(self._raw_package)
|
||||
except ValidationError as e:
|
||||
logger.error(f"Error validating package data: {e}")
|
||||
raise PackageImportException("Deserialization failed")
|
||||
|
||||
package_uuid = str(uuid.uuid4())
|
||||
if self.preserve_uuids:
|
||||
package_uuid = parsed.uuid
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# Import Elements
|
||||
self._import_compounds(package, parsed.compounds)
|
||||
self._import_simple_rules(package, parsed.simple_rules)
|
||||
self._import_composite_rules(package, parsed.composite_rules)
|
||||
self._import_reactions(package, parsed.reactions)
|
||||
self._import_pathways(package, parsed.pathways)
|
||||
self._import_scenarios(package, parsed.scenarios)
|
||||
self._import_additional_information(package, parsed.additional_information)
|
||||
self._link_scenarios_after_import(parsed)
|
||||
|
||||
return package
|
||||
|
||||
@staticmethod
|
||||
def sign(data: Dict[str, Any], key: str) -> Dict[str, Any]:
|
||||
@ -377,510 +640,6 @@ class PackageImporter:
|
||||
expected = hmac.new(key.encode(), json_str.encode(), hashlib.sha256).digest()
|
||||
return hmac.compare_digest(signature, expected)
|
||||
|
||||
@transaction.atomic
|
||||
def _import_package_from_json(self, package_data: Dict[str, Any]) -> Package:
|
||||
"""
|
||||
Import a complete package from JSON data.
|
||||
|
||||
Args:
|
||||
package_data: Dictionary containing the package export data
|
||||
|
||||
Returns:
|
||||
The created Package instance
|
||||
"""
|
||||
print(f"Starting import of package: {package_data['name']}")
|
||||
|
||||
# Create the main package
|
||||
package = self._create_package(package_data)
|
||||
|
||||
# Import in dependency order
|
||||
self._import_compounds(package, package_data.get("compounds", []))
|
||||
self._import_structures(package, package_data.get("structures", []))
|
||||
self._import_rules(package, package_data.get("rules", {}))
|
||||
self._import_reactions(package, package_data.get("reactions", []))
|
||||
self._import_pathways(package, package_data.get("pathways", []))
|
||||
self._import_nodes(package, package_data.get("nodes", []))
|
||||
self._import_edges(package, package_data.get("edges", []))
|
||||
self._import_scenarios(package, package_data.get("scenarios", []))
|
||||
|
||||
if package_data.get("models"):
|
||||
self._import_models(package, package_data["models"])
|
||||
|
||||
# Set default structures for compounds (after all structures are created)
|
||||
self._set_default_structures(package_data.get("compounds", []))
|
||||
|
||||
print(f"Package import completed: {package.name}")
|
||||
return package
|
||||
|
||||
def _create_package(self, package_data: Dict[str, Any]) -> Package:
|
||||
"""Create the main package object."""
|
||||
package_uuid = self._get_or_generate_uuid(package_data["uuid"])
|
||||
|
||||
# Handle license
|
||||
license_obj = None
|
||||
if package_data.get("license"):
|
||||
license_data = package_data["license"]
|
||||
license_obj, _ = License.objects.get_or_create(
|
||||
name=license_data["name"],
|
||||
defaults={
|
||||
"cc_string": license_data.get("cc_string", ""),
|
||||
"link": license_data.get("link", ""),
|
||||
"image_link": license_data.get("image_link", ""),
|
||||
},
|
||||
)
|
||||
|
||||
new_name = package_data.get("name")
|
||||
if self.add_import_timestamp:
|
||||
new_name = f"{new_name} - Imported at {datetime.now()}"
|
||||
|
||||
new_reviewed = False
|
||||
if self.trust_reviewed:
|
||||
new_reviewed = package_data.get("reviewed", False)
|
||||
|
||||
package = Package.objects.create(
|
||||
uuid=package_uuid,
|
||||
name=new_name,
|
||||
description=package_data["description"],
|
||||
kv=package_data.get("kv", {}),
|
||||
reviewed=new_reviewed,
|
||||
license=license_obj,
|
||||
)
|
||||
|
||||
self._cache_object("Package", package_data["uuid"], package)
|
||||
print(f"Created package: {package.name}")
|
||||
return package
|
||||
|
||||
def _import_compounds(self, package: Package, compounds_data: List[Dict[str, Any]]):
|
||||
"""Import compounds."""
|
||||
print(f"Importing {len(compounds_data)} compounds...")
|
||||
|
||||
for compound_data in compounds_data:
|
||||
compound_uuid = self._get_or_generate_uuid(compound_data["uuid"])
|
||||
|
||||
compound = Compound.objects.create(
|
||||
uuid=compound_uuid,
|
||||
package=package,
|
||||
name=compound_data["name"],
|
||||
description=compound_data["description"],
|
||||
kv=compound_data.get("kv", {}),
|
||||
# default_structure will be set later
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if compound_data.get("aliases"):
|
||||
compound.aliases = compound_data["aliases"]
|
||||
compound.save()
|
||||
|
||||
self._cache_object("Compound", compound_data["uuid"], compound)
|
||||
|
||||
# Handle external identifiers
|
||||
self._create_external_identifiers(
|
||||
compound, compound_data.get("external_identifiers", [])
|
||||
)
|
||||
|
||||
def _import_structures(self, package: Package, structures_data: List[Dict[str, Any]]):
|
||||
"""Import compound structures."""
|
||||
print(f"Importing {len(structures_data)} compound structures...")
|
||||
|
||||
for structure_data in structures_data:
|
||||
structure_uuid = self._get_or_generate_uuid(structure_data["uuid"])
|
||||
compound_uuid = structure_data["compound"]["uuid"]
|
||||
compound = self._get_cached_object("Compound", compound_uuid)
|
||||
|
||||
if not compound:
|
||||
print(f"Warning: Compound with UUID {compound_uuid} not found for structure")
|
||||
continue
|
||||
|
||||
structure = CompoundStructure.objects.create(
|
||||
uuid=structure_uuid,
|
||||
compound=compound,
|
||||
name=structure_data["name"],
|
||||
description=structure_data["description"],
|
||||
kv=structure_data.get("kv", {}),
|
||||
smiles=structure_data["smiles"],
|
||||
canonical_smiles=structure_data["canonical_smiles"],
|
||||
inchikey=structure_data["inchikey"],
|
||||
normalized_structure=structure_data.get("normalized_structure", False),
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if structure_data.get("aliases"):
|
||||
structure.aliases = structure_data["aliases"]
|
||||
structure.save()
|
||||
|
||||
self._cache_object("CompoundStructure", structure_data["uuid"], structure)
|
||||
|
||||
# Handle external identifiers
|
||||
self._create_external_identifiers(
|
||||
structure, structure_data.get("external_identifiers", [])
|
||||
)
|
||||
|
||||
def _import_rules(self, package: Package, rules_data: Dict[str, Any]):
|
||||
"""Import all types of rules."""
|
||||
print("Importing rules...")
|
||||
|
||||
# Import simple rules first
|
||||
simple_rules_data = rules_data.get("simple_rules", [])
|
||||
print(f"Importing {len(simple_rules_data)} simple rules...")
|
||||
|
||||
for rule_data in simple_rules_data:
|
||||
self._create_simple_rule(package, rule_data)
|
||||
|
||||
# Import parallel rules
|
||||
parallel_rules_data = rules_data.get("parallel_rules", [])
|
||||
print(f"Importing {len(parallel_rules_data)} parallel rules...")
|
||||
|
||||
for rule_data in parallel_rules_data:
|
||||
self._create_parallel_rule(package, rule_data)
|
||||
|
||||
def _create_simple_rule(self, package: Package, rule_data: Dict[str, Any]):
|
||||
"""Create a simple rule (SimpleAmbitRule or SimpleRDKitRule)."""
|
||||
rule_uuid = self._get_or_generate_uuid(rule_data["uuid"])
|
||||
rule_type = rule_data.get("rule_type", "SimpleRule")
|
||||
|
||||
common_fields = {
|
||||
"uuid": rule_uuid,
|
||||
"package": package,
|
||||
"name": rule_data["name"],
|
||||
"description": rule_data["description"],
|
||||
"kv": rule_data.get("kv", {}),
|
||||
}
|
||||
|
||||
if rule_type == "SimpleAmbitRule":
|
||||
rule = SimpleAmbitRule.objects.create(
|
||||
**common_fields,
|
||||
smirks=rule_data.get("smirks", ""),
|
||||
reactant_filter_smarts=rule_data.get("reactant_filter_smarts", ""),
|
||||
product_filter_smarts=rule_data.get("product_filter_smarts", ""),
|
||||
)
|
||||
elif rule_type == "SimpleRDKitRule":
|
||||
rule = SimpleRDKitRule.objects.create(
|
||||
**common_fields, reaction_smarts=rule_data.get("reaction_smarts", "")
|
||||
)
|
||||
else:
|
||||
rule = SimpleRule.objects.create(**common_fields)
|
||||
|
||||
# Set aliases if present
|
||||
if rule_data.get("aliases"):
|
||||
rule.aliases = rule_data["aliases"]
|
||||
rule.save()
|
||||
|
||||
self._cache_object("SimpleRule", rule_data["uuid"], rule)
|
||||
return rule
|
||||
|
||||
def _create_parallel_rule(self, package: Package, rule_data: Dict[str, Any]):
|
||||
"""Create a parallel rule."""
|
||||
rule_uuid = self._get_or_generate_uuid(rule_data["uuid"])
|
||||
|
||||
rule = ParallelRule.objects.create(
|
||||
uuid=rule_uuid,
|
||||
package=package,
|
||||
name=rule_data["name"],
|
||||
description=rule_data["description"],
|
||||
kv=rule_data.get("kv", {}),
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if rule_data.get("aliases"):
|
||||
rule.aliases = rule_data["aliases"]
|
||||
rule.save()
|
||||
|
||||
# Add simple rules
|
||||
for simple_rule_ref in rule_data.get("simple_rules", []):
|
||||
simple_rule = self._get_cached_object("SimpleRule", simple_rule_ref["uuid"])
|
||||
if simple_rule:
|
||||
rule.simple_rules.add(simple_rule)
|
||||
|
||||
self._cache_object("ParallelRule", rule_data["uuid"], rule)
|
||||
return rule
|
||||
|
||||
def _import_reactions(self, package: Package, reactions_data: List[Dict[str, Any]]):
|
||||
"""Import reactions."""
|
||||
print(f"Importing {len(reactions_data)} reactions...")
|
||||
|
||||
for reaction_data in reactions_data:
|
||||
reaction_uuid = self._get_or_generate_uuid(reaction_data["uuid"])
|
||||
|
||||
reaction = Reaction.objects.create(
|
||||
uuid=reaction_uuid,
|
||||
package=package,
|
||||
name=reaction_data["name"],
|
||||
description=reaction_data["description"],
|
||||
kv=reaction_data.get("kv", {}),
|
||||
multi_step=reaction_data.get("multi_step", False),
|
||||
medline_references=reaction_data.get("medline_references", []),
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if reaction_data.get("aliases"):
|
||||
reaction.aliases = reaction_data["aliases"]
|
||||
reaction.save()
|
||||
|
||||
# Add educts and products
|
||||
for educt_ref in reaction_data.get("educts", []):
|
||||
compound = self._get_cached_object("CompoundStructure", educt_ref["uuid"])
|
||||
if compound:
|
||||
reaction.educts.add(compound)
|
||||
|
||||
for product_ref in reaction_data.get("products", []):
|
||||
compound = self._get_cached_object("CompoundStructure", product_ref["uuid"])
|
||||
if compound:
|
||||
reaction.products.add(compound)
|
||||
|
||||
# Add rules
|
||||
for rule_ref in reaction_data.get("rules", []):
|
||||
# Try to find rule in different caches
|
||||
rule = self._get_cached_object(
|
||||
"SimpleRule", rule_ref["uuid"]
|
||||
) or self._get_cached_object("ParallelRule", rule_ref["uuid"])
|
||||
if rule:
|
||||
reaction.rules.add(rule)
|
||||
|
||||
self._cache_object("Reaction", reaction_data["uuid"], reaction)
|
||||
|
||||
# Handle external identifiers
|
||||
self._create_external_identifiers(
|
||||
reaction, reaction_data.get("external_identifiers", [])
|
||||
)
|
||||
|
||||
def _import_pathways(self, package: Package, pathways_data: List[Dict[str, Any]]):
|
||||
"""Import pathways."""
|
||||
print(f"Importing {len(pathways_data)} pathways...")
|
||||
|
||||
for pathway_data in pathways_data:
|
||||
pathway_uuid = self._get_or_generate_uuid(pathway_data["uuid"])
|
||||
|
||||
pathway = Pathway.objects.create(
|
||||
uuid=pathway_uuid,
|
||||
package=package,
|
||||
name=pathway_data["name"],
|
||||
description=pathway_data["description"],
|
||||
kv=pathway_data.get("kv", {}),
|
||||
# setting will be handled separately if needed
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if pathway_data.get("aliases"):
|
||||
pathway.aliases = pathway_data["aliases"]
|
||||
pathway.save()
|
||||
|
||||
self._cache_object("Pathway", pathway_data["uuid"], pathway)
|
||||
|
||||
def _import_nodes(self, package: Package, nodes_data: List[Dict[str, Any]]):
|
||||
"""Import pathway nodes."""
|
||||
print(f"Importing {len(nodes_data)} nodes...")
|
||||
|
||||
for node_data in nodes_data:
|
||||
node_uuid = self._get_or_generate_uuid(node_data["uuid"])
|
||||
pathway_uuid = node_data["pathway"]["uuid"]
|
||||
pathway = self._get_cached_object("Pathway", pathway_uuid)
|
||||
|
||||
if not pathway:
|
||||
print(f"Warning: Pathway with UUID {pathway_uuid} not found for node")
|
||||
continue
|
||||
|
||||
# For now, we'll set default_node_label to None and handle it later
|
||||
# as it requires compound structures to be fully imported
|
||||
node = Node.objects.create(
|
||||
uuid=node_uuid,
|
||||
pathway=pathway,
|
||||
name=node_data["name"],
|
||||
description=node_data["description"],
|
||||
kv=node_data.get("kv", {}),
|
||||
depth=node_data.get("depth", 0),
|
||||
default_node_label=self._get_cached_object(
|
||||
"CompoundStructure", node_data["default_node_label"]["uuid"]
|
||||
),
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if node_data.get("aliases"):
|
||||
node.aliases = node_data["aliases"]
|
||||
node.save()
|
||||
|
||||
self._cache_object("Node", node_data["uuid"], node)
|
||||
# Store node_data for later processing of relationships
|
||||
node._import_data = node_data
|
||||
|
||||
def _import_edges(self, package: Package, edges_data: List[Dict[str, Any]]):
|
||||
"""Import pathway edges."""
|
||||
print(f"Importing {len(edges_data)} edges...")
|
||||
|
||||
for edge_data in edges_data:
|
||||
edge_uuid = self._get_or_generate_uuid(edge_data["uuid"])
|
||||
pathway_uuid = edge_data["pathway"]["uuid"]
|
||||
pathway = self._get_cached_object("Pathway", pathway_uuid)
|
||||
|
||||
if not pathway:
|
||||
print(f"Warning: Pathway with UUID {pathway_uuid} not found for edge")
|
||||
continue
|
||||
|
||||
# For now, we'll set edge_label to None and handle it later
|
||||
edge = Edge.objects.create(
|
||||
uuid=edge_uuid,
|
||||
pathway=pathway,
|
||||
name=edge_data["name"],
|
||||
description=edge_data["description"],
|
||||
kv=edge_data.get("kv", {}),
|
||||
edge_label=self._get_cached_object("Reaction", edge_data["edge_label"]["uuid"]),
|
||||
)
|
||||
|
||||
# Set aliases if present
|
||||
if edge_data.get("aliases"):
|
||||
edge.aliases = edge_data["aliases"]
|
||||
edge.save()
|
||||
|
||||
# Add start and end nodes
|
||||
for start_node_ref in edge_data.get("start_nodes", []):
|
||||
node = self._get_cached_object("Node", start_node_ref["uuid"])
|
||||
if node:
|
||||
edge.start_nodes.add(node)
|
||||
|
||||
for end_node_ref in edge_data.get("end_nodes", []):
|
||||
node = self._get_cached_object("Node", end_node_ref["uuid"])
|
||||
if node:
|
||||
edge.end_nodes.add(node)
|
||||
|
||||
self._cache_object("Edge", edge_data["uuid"], edge)
|
||||
|
||||
def _import_scenarios(self, package: Package, scenarios_data: List[Dict[str, Any]]):
|
||||
"""Import scenarios."""
|
||||
print(f"Importing {len(scenarios_data)} scenarios...")
|
||||
|
||||
# First pass: create scenarios without parent relationships
|
||||
for scenario_data in scenarios_data:
|
||||
scenario_uuid = self._get_or_generate_uuid(scenario_data["uuid"])
|
||||
|
||||
scenario_date = None
|
||||
if scenario_data.get("scenario_date"):
|
||||
scenario_date = scenario_data["scenario_date"]
|
||||
|
||||
scenario = Scenario.objects.create(
|
||||
uuid=scenario_uuid,
|
||||
package=package,
|
||||
name=scenario_data["name"],
|
||||
description=scenario_data["description"],
|
||||
kv=scenario_data.get("kv", {}),
|
||||
scenario_date=scenario_date,
|
||||
scenario_type=scenario_data.get("scenario_type"),
|
||||
additional_information=scenario_data.get("additional_information", {}),
|
||||
)
|
||||
|
||||
self._cache_object("Scenario", scenario_data["uuid"], scenario)
|
||||
# Store scenario_data for later processing of parent relationships
|
||||
scenario._import_data = scenario_data
|
||||
|
||||
# Second pass: set parent relationships
|
||||
for scenario_data in scenarios_data:
|
||||
if scenario_data.get("parent"):
|
||||
scenario = self._get_cached_object("Scenario", scenario_data["uuid"])
|
||||
parent = self._get_cached_object("Scenario", scenario_data["parent"]["uuid"])
|
||||
if scenario and parent:
|
||||
scenario.parent = parent
|
||||
scenario.save()
|
||||
|
||||
def _import_models(self, package: Package, models_data: List[Dict[str, Any]]):
|
||||
"""Import EPModels."""
|
||||
print(f"Importing {len(models_data)} models...")
|
||||
|
||||
for model_data in models_data:
|
||||
model_uuid = self._get_or_generate_uuid(model_data["uuid"])
|
||||
model_type = model_data.get("model_type", "EPModel")
|
||||
|
||||
common_fields = {
|
||||
"uuid": model_uuid,
|
||||
"package": package,
|
||||
"name": model_data["name"],
|
||||
"description": model_data["description"],
|
||||
"kv": model_data.get("kv", {}),
|
||||
}
|
||||
|
||||
# Add PackageBasedModel fields if present
|
||||
if "threshold" in model_data:
|
||||
common_fields.update(
|
||||
{
|
||||
"threshold": model_data.get("threshold"),
|
||||
"eval_results": model_data.get("eval_results", {}),
|
||||
"model_status": model_data.get("model_status", "INITIAL"),
|
||||
}
|
||||
)
|
||||
|
||||
# Create the appropriate model type
|
||||
if model_type == "RuleBasedRelativeReasoning":
|
||||
model = RuleBasedRelativeReasoning.objects.create(
|
||||
**common_fields,
|
||||
min_count=model_data.get("min_count", 1),
|
||||
max_count=model_data.get("max_count", 10),
|
||||
)
|
||||
elif model_type == "MLRelativeReasoning":
|
||||
model = MLRelativeReasoning.objects.create(**common_fields)
|
||||
elif model_type == "EnviFormer":
|
||||
model = EnviFormer.objects.create(**common_fields)
|
||||
else:
|
||||
model = EPModel.objects.create(**common_fields)
|
||||
|
||||
# Set aliases if present
|
||||
if model_data.get("aliases"):
|
||||
model.aliases = model_data["aliases"]
|
||||
model.save()
|
||||
|
||||
# Add package relationships for PackageBasedModel
|
||||
if hasattr(model, "rule_packages"):
|
||||
for pkg_ref in model_data.get("rule_packages", []):
|
||||
pkg = self._get_cached_object("Package", pkg_ref["uuid"])
|
||||
if pkg:
|
||||
model.rule_packages.add(pkg)
|
||||
|
||||
for pkg_ref in model_data.get("data_packages", []):
|
||||
pkg = self._get_cached_object("Package", pkg_ref["uuid"])
|
||||
if pkg:
|
||||
model.data_packages.add(pkg)
|
||||
|
||||
for pkg_ref in model_data.get("eval_packages", []):
|
||||
pkg = self._get_cached_object("Package", pkg_ref["uuid"])
|
||||
if pkg:
|
||||
model.eval_packages.add(pkg)
|
||||
|
||||
self._cache_object("EPModel", model_data["uuid"], model)
|
||||
|
||||
def _set_default_structures(self, compounds_data: List[Dict[str, Any]]):
|
||||
"""Set default structures for compounds after all structures are created."""
|
||||
print("Setting default structures for compounds...")
|
||||
|
||||
for compound_data in compounds_data:
|
||||
if compound_data.get("default_structure"):
|
||||
compound = self._get_cached_object("Compound", compound_data["uuid"])
|
||||
structure = self._get_cached_object(
|
||||
"CompoundStructure", compound_data["default_structure"]["uuid"]
|
||||
)
|
||||
if compound and structure:
|
||||
compound.default_structure = structure
|
||||
compound.save()
|
||||
|
||||
def _create_external_identifiers(self, obj, identifiers_data: List[Dict[str, Any]]):
|
||||
"""Create external identifiers for an object."""
|
||||
for identifier_data in identifiers_data:
|
||||
# Get or create the external database
|
||||
db_data = identifier_data["database"]
|
||||
database, _ = ExternalDatabase.objects.get_or_create(
|
||||
name=db_data["name"],
|
||||
defaults={
|
||||
"base_url": db_data.get("base_url", ""),
|
||||
"full_name": db_data.get("name", ""),
|
||||
"description": "",
|
||||
"is_active": True,
|
||||
},
|
||||
)
|
||||
|
||||
# Create the external identifier
|
||||
ExternalIdentifier.objects.create(
|
||||
content_object=obj,
|
||||
database=database,
|
||||
identifier_value=identifier_data["identifier_value"],
|
||||
url=identifier_data.get("url", ""),
|
||||
is_primary=identifier_data.get("is_primary", False),
|
||||
)
|
||||
|
||||
|
||||
class PathwayUtils:
|
||||
def __init__(self, pathway: "Pathway"):
|
||||
|
||||
Reference in New Issue
Block a user