Added new swagger model and other small changes

mild-blue · Dec 7, 2022 · 08f76d4 · 08f76d4
1 parent fd3b915
commit 08f76d4
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 14 deletions.
diff --git a/txmatching/data_transfer_objects/optimizer/optimizer_in_swagger.py b/txmatching/data_transfer_objects/optimizer/optimizer_in_swagger.py
@@ -52,14 +52,16 @@ def output(self, key, obj, **kwargs):
                                              [{"hla_compatibility_score": 3}, {"donor_age_difference": 20}]])
 })
 
+CompGraphEntry = optimizer_api.model("CompatibilityGraphEntry", {
+    "donor_index": fields.Integer(required=True, example=1),
+    "recipient_index": fields.Integer(required=True, example=2),
+    "weights": DictItem(attribute="calling_args", example={"hla_compatibility_score": 17,
+                                                           "donor_age_difference": 1
+                                                           })
+})
+
 OptimizerRequestObjectJson = optimizer_api.model('OptimizerRequest', {
-    'compatibility_graph': fields.List(required=True, cls_or_instance=DictItem(attribute="calling_args",
-                                                                               example={"donor_index": 1,
-                                                                                        "recipient_index": 2,
-                                                                                        "weights": {
-                                                                                            "hla_compatibility_score": 17,
-                                                                                            "donor_age_difference": 1
-                                                                                        }})),
+    'compatibility_graph': fields.List(required=True, cls_or_instance=fields.Nested(CompGraphEntry)),
     'pairs': fields.List(reqired=True, cls_or_instance=fields.Nested(PairJson)),
     'configuration': fields.Nested(OptimizerConfigurationJson, required=True)
 })
diff --git a/txmatching/optimizer/optimizer_functions.py b/txmatching/optimizer/optimizer_functions.py
@@ -4,11 +4,16 @@
 import numpy as np
 
 from txmatching.configuration.config_parameters import ConfigParameters
-from txmatching.optimizer.optimizer_return_object import CycleOrChain, DonorToRecipient, OptimizerReturn, Statistics
-from txmatching.optimizer.optimizer_request_object import CompatibilityGraphEntry, Limitations, OptimizerConfiguration, \
-    OptimizerRequest, Pair
-from txmatching.patients.patient_types import DonorDbId, RecipientDbId
+from txmatching.optimizer.optimizer_request_object import (
+    CompatibilityGraphEntry, Limitations, OptimizerConfiguration,
+    OptimizerRequest, Pair)
+from txmatching.optimizer.optimizer_return_object import (CycleOrChain,
+                                                          DonorToRecipient,
+                                                          OptimizerReturn,
+                                                          Statistics)
 from txmatching.patients.patient import Donor, Recipient
+from txmatching.patients.patient_types import DonorDbId, RecipientDbId
+from txmatching.scorers.compatibility_graph import OptimizerCompatibilityGraph
 from txmatching.scorers.scorer_from_config import scorer_from_configuration
 from txmatching.scorers.split_hla_additive_scorer import SplitScorer
 from txmatching.solve_service.solver_lock import run_with_solver_lock
@@ -17,7 +22,6 @@
 from txmatching.solvers.solver_from_config import solver_from_configuration
 from txmatching.utils.blood_groups import BloodGroup
 from txmatching.utils.enums import Solver
-from txmatching.scorers.compatibility_graph import OptimizerCompatibilityGraph
 
 
 def calculate_from_optimizer_safe(optimizer_request_object: OptimizerRequest) -> OptimizerReturn:

diff --git a/txmatching/scorers/compatibility_graph.py b/txmatching/scorers/compatibility_graph.py
@@ -1,5 +1,4 @@
-from typing import Dict, Tuple, List
-
+from typing import Dict, List, Tuple
 
 CompatibilityGraph = Dict[Tuple[int, int], Dict[str, int]]