Improved transmission distance creation

ensysmod/crud/base_depends_component.py

@@ -13,6 +13,9 @@ class CRUDBaseDependsComponent(CRUDBaseDependsDataset, Generic[ModelType, Create
     Base class for all CRUD classes that depend on a dataset and component.
     """
 
+    def get(self, db: Session, component_id: int) -> Optional[ModelType]:
+        return db.query(self.model).filter(self.model.ref_component == component_id).first()
+
     def get_by_dataset_and_name(self, db: Session, *, dataset_id: int, name: str) -> Optional[ModelType]:
         """
         Get an energy component based object by dataset and name.

ensysmod/crud/energy_transmission.py

@@ -23,6 +23,7 @@ def create(self, db: Session, *, obj_in: EnergyTransmissionCreate) -> EnergyTran
 
         # also create distances
         for distance_create in obj_in.distances:
+            distance_create.ref_dataset = obj_in.ref_dataset
             distance_create.ref_component = db_obj.component.id
             crud.energy_transmission_distance.create(db, obj_in=distance_create)
 

ensysmod/crud/energy_transmission_distance.py

@@ -1,5 +1,6 @@
 from sqlalchemy.orm import Session
 
+from ensysmod import crud
 from ensysmod.crud.base import CRUDBase
 from ensysmod.model import EnergyTransmissionDistance
 from ensysmod.schemas import EnergyTransmissionDistanceCreate, EnergyTransmissionDistanceUpdate
@@ -22,5 +23,51 @@ def remove_by_component(self, db: Session, component_id: int):
         """
         db.query(EnergyTransmissionDistance).filter(EnergyTransmissionDistance.ref_component == component_id).delete()
 
+    def create(self, db: Session, obj_in: EnergyTransmissionDistanceCreate) -> EnergyTransmissionDistance:
+        """
+        Creates a new energy transmission distance entry between two regions.
+
+        :param db: Database session
+        :param obj_in: Input data
+        :return: New energy transmission distance entry
+        """
+
+        if obj_in.ref_component is None and obj_in.component is None:
+            raise ValueError("Component must be specified. Provide reference id or component name.")
+
+        if obj_in.ref_region_from is None and obj_in.region_from is None:
+            raise ValueError("Region from must be specified. Provide reference id or region name.")
+
+        if obj_in.ref_component is not None:
+            transmission = crud.energy_transmission.get(db, obj_in.ref_component)
+        else:
+            transmission = crud.energy_transmission.get_by_dataset_and_name(db, dataset_id=obj_in.ref_dataset,
+                                                                            name=obj_in.component)
+
+        if transmission is None or transmission.component.ref_dataset != obj_in.ref_dataset:
+            raise ValueError("Component not found or from different dataset.")
+        obj_in.ref_component = transmission.ref_component
+
+        if obj_in.ref_region_from is not None:
+            region_from = crud.region.get(db, obj_in.ref_region_from)
+        else:
+            region_from = crud.region.get_by_dataset_and_name(db, dataset_id=obj_in.ref_dataset,
+                                                              name=obj_in.region_from)
+
+        if region_from is None or region_from.ref_dataset != obj_in.ref_dataset:
+            raise ValueError("Region from not found or from different dataset.")
+        obj_in.ref_region_from = region_from.id
+
+        if obj_in.ref_region_to is not None:
+            region_to = crud.region.get(db, obj_in.ref_region_to)
+        else:
+            region_to = crud.region.get_by_dataset_and_name(db, dataset_id=obj_in.ref_dataset, name=obj_in.region_to)
+
+        if region_to is None or region_to.ref_dataset != obj_in.ref_dataset:
+            raise ValueError("Region to not found or from different dataset.")
+        obj_in.ref_region_to = region_to.id
+
+        return super().create(db=db, obj_in=obj_in)
+
 
 energy_transmission_distance = CRUDEnergyTransmissionDistance(EnergyTransmissionDistance)

ensysmod/model/energy_transmission_distance.py

@@ -10,8 +10,8 @@ class EnergyTransmissionDistance(Base):
 
     See https://vsa-fine.readthedocs.io/en/latest/storageClassDoc.html
     """
-    ref_component = Column(Integer, ForeignKey("energy_transmission.ref_component"), index=True, nullable=False,
-                           primary_key=True)
+    id = Column(Integer, primary_key=True)
+    ref_component = Column(Integer, ForeignKey("energy_transmission.ref_component"), index=True, nullable=False)
     ref_region_from = Column(Integer, ForeignKey("region.id"), index=True, nullable=False)
     ref_region_to = Column(Integer, ForeignKey("region.id"), index=True, nullable=False)
 

ensysmod/schemas/energy_transmission_distance.py

@@ -17,9 +17,15 @@ class EnergyTransmissionDistanceCreate(EnergyTransmissionDistanceBase):
     Properties to receive via API on creation of an energy transmission distance.
     """
     ref_dataset: Optional[int] = None
+
     ref_component: Optional[int] = None
-    ref_region_from: int
-    ref_region_to: int
+    component: Optional[str] = None
+
+    ref_region_from: Optional[int] = None
+    region_from: Optional[str] = None
+
+    ref_region_to: Optional[int] = None
+    region_to: Optional[str] = None
 
 
 class EnergyTransmissionDistanceUpdate(EnergyTransmissionDistanceBase):
@@ -33,6 +39,7 @@ class EnergyTransmissionDistance(EnergyTransmissionDistanceBase):
     """
     Properties to return via API for an energy transmission distance.
     """
+    id: int
     region_from: Region
     region_to: Region
 

tests/utils/data_generator/__init__.py

@@ -10,4 +10,5 @@
     random_energy_storage_create
 from .energy_transmissions import random_existing_energy_transmission, fixed_existing_energy_transmission, \
     random_energy_transmission_create
-from .regions import random_existing_region, fixed_existing_region, random_region_create
+from .regions import random_existing_region, fixed_existing_region, random_region_create, \
+    fixed_alternative_existing_region

tests/utils/data_generator/energy_transmissions.py

@@ -5,15 +5,17 @@
 from ensysmod.schemas import EnergyTransmissionCreate
 from ensysmod.schemas.energy_transmission_distance import EnergyTransmissionDistanceCreate
 from tests.utils.data_generator import fixed_existing_dataset, fixed_existing_energy_commodity
-from tests.utils.data_generator.regions import random_existing_region, fixed_existing_region
+from tests.utils.data_generator.regions import fixed_existing_region, fixed_alternative_existing_region, \
+    fixed_alternative_alternative_existing_region
 from tests.utils.utils import random_lower_string
 
 
 def random_energy_transmission_create(db: Session) -> EnergyTransmissionCreate:
     dataset = fixed_existing_dataset(db)
     commodity = fixed_existing_energy_commodity(db)
-    region = random_existing_region(db)
-    region_to = random_existing_region(db)
+    region = fixed_existing_region(db)
+    region_to = fixed_alternative_existing_region(db)
+    region_to_alt = fixed_alternative_alternative_existing_region(db)
     return EnergyTransmissionCreate(
         ref_dataset=dataset.id,
         name=f"EnergyTransmission-{dataset.id}-{random_lower_string()}",
@@ -24,7 +26,12 @@ def random_energy_transmission_create(db: Session) -> EnergyTransmissionCreate:
             EnergyTransmissionDistanceCreate(
                 distance=42.3,
                 ref_region_from=region.id,
-                ref_region_to=region_to.id,
+                region_to=region_to.name,
+            ),
+            EnergyTransmissionDistanceCreate(
+                distance=44.3,
+                region_from=region.name,
+                ref_region_to=region_to_alt.id,
             )
         ]
     )

tests/utils/data_generator/regions.py

@@ -24,9 +24,37 @@ def fixed_region_create(db: Session) -> RegionCreate:
                         ref_dataset=dataset.id)
 
 
+def fixed_alternative_region_create(db: Session) -> RegionCreate:
+    dataset = fixed_existing_dataset(db)
+    return RegionCreate(name=f"Region-{dataset.id}-Fixed-alternative",
+                        ref_dataset=dataset.id)
+
+
+def fixed_alternative_alternative_region_create(db: Session) -> RegionCreate:
+    dataset = fixed_existing_dataset(db)
+    return RegionCreate(name=f"Region-{dataset.id}-Fixed-alternative-alternative",
+                        ref_dataset=dataset.id)
+
+
 def fixed_existing_region(db: Session) -> Region:
     create_request = fixed_region_create(db)
     region = crud.region.get_by_dataset_and_name(db=db, dataset_id=create_request.ref_dataset, name=create_request.name)
     if region is None:
         region = crud.region.create(db=db, obj_in=create_request)
     return region
+
+
+def fixed_alternative_existing_region(db: Session) -> Region:
+    create_request = fixed_alternative_region_create(db)
+    region = crud.region.get_by_dataset_and_name(db=db, dataset_id=create_request.ref_dataset, name=create_request.name)
+    if region is None:
+        region = crud.region.create(db=db, obj_in=create_request)
+    return region
+
+
+def fixed_alternative_alternative_existing_region(db: Session) -> Region:
+    create_request = fixed_alternative_alternative_region_create(db)
+    region = crud.region.get_by_dataset_and_name(db=db, dataset_id=create_request.ref_dataset, name=create_request.name)
+    if region is None:
+        region = crud.region.create(db=db, obj_in=create_request)
+    return region