Added Plan Status Code Capabilities to PyAres

PyAres/Models/__init__.py

@@ -1,4 +1,4 @@
-from .ares_data_models import AresDataType, Outcome, RequestMetadata, AresSchemaEntry, Quantity, QuantitySchema, Limits
+from .ares_data_models import AresDataType, Outcome, RequestMetadata, AresSchemaEntry, Quantity, QuantitySchema, Limits, PlanStatusCode
 
 __all__ = [
     "AresDataType",
@@ -7,5 +7,6 @@
     "AresSchemaEntry",
     "Quantity",
     "QuantitySchema",
-    "Limits"
+    "Limits",
+    "PlanStatusCode"
 ]

PyAres/Models/ares_data_models.py

@@ -30,6 +30,12 @@ class Outcome(Enum):
     WARNING = 3
     CANCELED = 4
 
+class PlanStatusCode(Enum):
+    PLAN_STATUS_UNSPECIFIED = 0
+    PLAN_ACCEPTED = 1
+    PLAN_UNACHIEVABLE = 2
+    PLAN_FAILED = 3
+
 class RequestMetadata():
     def __init__(self, proto_metadata: request_metadata_pb2.RequestMetadata):
         self.system_name = proto_metadata.system_name

PyAres/Planning/planner_models.py

@@ -1,5 +1,6 @@
 from typing import Dict, Any, List, Sequence, Optional
-from ..Models import Outcome, AresDataType, RequestMetadata
+from ..Models import Outcome, AresDataType, RequestMetadata, PlanStatusCode
+from enum import Enum
 
 class ParameterHistoryItem:
     """ Represents a single historical parameter item """
@@ -107,7 +108,12 @@ class PlanRequest:
 
     Designed to provide a more user-friendly abstraction for interacting with a plan request message.
     """
-    def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any], analysis_results: Sequence[float], metadata: RequestMetadata = RequestMetadata.from_default_values()):
+    def __init__(self, 
+                 parameters: list[PlanningParameter], 
+                 settings: Dict[str, Any], 
+                 analysis_results: Sequence[float], 
+                 metadata: RequestMetadata = RequestMetadata.from_default_values(),
+                 previous_plan_status_code: PlanStatusCode = PlanStatusCode.PLAN_STATUS_UNSPECIFIED):
         """
         Initializes a PlanRequest.
 
@@ -118,6 +124,7 @@ def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any]
         self.settings = settings
         self.analysis_results = analysis_results
         self.request_metadata = metadata
+        self.previous_plan_status_code = previous_plan_status_code
 
     def __str__(self) -> str:
         param_str = "\n ".join(self.parameter_names)
@@ -175,6 +182,8 @@ def __init__(self,
             parameter_names: A list of names associated with planned parameters.
             parameter_values: A list of values associated with planned parameters. 
             parameter_data: A python dictionary of key:value pairs of planned parameters and planned values
+            outcome: An enum of type Outcome that determines whether the planning process succeeded or not, defaults to SUCCESS
+            error_string: An optional string for specifying planning failure reasons to be relayed to ARES
         """
         if parameter_data is not None:
             self.parameter_names = list(parameter_data.keys())

PyAres/Planning/planning_service.py

@@ -21,6 +21,7 @@
 from ..Utils import ares_data_type_utils
 from ..Utils import ares_struct_utils
 from ..Utils import ares_outcome_utils
+from ..Utils import ares_plan_status_code_utils
 
 # Import python models
 from ..Models import ares_data_models, Limits
@@ -122,7 +123,8 @@ def Plan(self, request: plan_pb2.PlanningRequest, context) -> plan_pb2.PlanningR
         python_request = PlanRequest(parameters=parameters, 
                                      settings=ares_struct_utils.ares_struct_to_dict(request.adapter_settings), 
                                      analysis_results=list(request.analysis_results),
-                                     metadata=RequestMetadata(request.metadata))
+                                     metadata=RequestMetadata(request.metadata),
+                                     previous_plan_status_code=ares_plan_status_code_utils.proto_plan_status_to_python_plan_status(request.previous_plan_status_code))
 
         #Handle call using the user's custom planning logic 
         response_proto = plan_pb2.PlanningResponse()

PyAres/Utils/ares_plan_status_code_utils.py

@@ -0,0 +1,10 @@
+from ..Models import PlanStatusCode
+from ares_datamodel.planning import plan_pb2
+from typing import cast
+
+def python_plan_status_to_proto_plan_status(py_value: PlanStatusCode) -> plan_pb2.PlanStatusCode:
+  val = cast(plan_pb2.PlanStatusCode, py_value)
+  return val
+
+def proto_plan_status_to_python_plan_status(proto_value: plan_pb2.PlanStatusCode) -> PlanStatusCode:
+  return PlanStatusCode(proto_value)