Add Protocol Errors

gufe/protocols/__init__.py

@@ -1,5 +1,13 @@
 """Defining processes for performing estimates of free energy differences"""
 
+from .errors import (
+    MissingProtocolUnitError,
+    ProtocolAnalysisError,
+    ProtocolExecutionError,
+    ProtocolSetupError,
+    ProtocolUnitFailureError,
+    ProtocolValidationError,
+)
 from .protocol import Protocol, ProtocolResult
 from .protocoldag import ProtocolDAG, ProtocolDAGResult, execute_DAG
 from .protocolunit import Context, ProtocolUnit, ProtocolUnitFailure, ProtocolUnitResult

gufe/protocols/errors.py

@@ -0,0 +1,28 @@
+# This code is part of OpenFE and is licensed under the MIT license.
+# For details, see https://github.com/OpenFreeEnergy/gufe
+
+
+# Protocol Errors
+class ProtocolValidationError(Exception):
+    """Error when the protocol setup or settings can not be validated."""
+
+
+class ProtocolSetupError(Exception):
+    """Error when executing the setup unit of the protocol."""
+
+
+class ProtocolExecutionError(Exception):
+    """Error when executing the production unit of the protocol."""
+
+
+class ProtocolAnalysisError(Exception):
+    """Error when trying to perform some analyses after the protocol has been executed."""
+
+
+# Protocol Results Errors
+class MissingProtocolUnitError(Exception):
+    """Error when a ProtocolDAGResult is massing a protocol unit."""
+
+
+class ProtocolUnitFailureError(Exception):
+    """Error when a ProtocolDAGResult contains an failed protocol unit."""

gufe/protocols/protocoldag.py

@@ -14,6 +14,7 @@
 import networkx as nx
 
 from ..tokenization import GufeKey, GufeTokenizable
+from .errors import MissingProtocolUnitError, ProtocolUnitFailureError
 from .protocolunit import Context, ProtocolUnit, ProtocolUnitFailure, ProtocolUnitResult
 
 
@@ -211,22 +212,24 @@ def unit_to_result(self, protocol_unit: ProtocolUnit) -> ProtocolUnitResult:
 
         Raises
         ------
-        KeyError
-          if either there are no results, or only failures
+        MissingProtocolUnitError:
+          if there are no results for that protocol unit
+        ProtocolUnitFailureError:
+          if all units failed
         """
         try:
             units = self._unit_result_mapping[protocol_unit]
         except KeyError:
-            raise KeyError("No such `protocol_unit` present")
+            raise MissingProtocolUnitError(f"No such `protocol_unit`:{protocol_unit} present")
         else:
             for u in units:
                 if u.ok():
                     return u
             else:
-                raise KeyError("No success for `protocol_unit` found")
+                raise ProtocolUnitFailureError(f"No success for `protocol_unit`:{protocol_unit} found")
 
     def unit_to_all_results(self, protocol_unit: ProtocolUnit) -> list[ProtocolUnitResult]:
-        """Return all results (sucess and failure) for a given Unit.
+        """Return all results (success and failure) for a given Unit.
 
         Returns
         -------
@@ -235,19 +238,19 @@ def unit_to_all_results(self, protocol_unit: ProtocolUnit) -> list[ProtocolUnitR
 
         Raises
         ------
-        KeyError
+        MissingProtocolUnitError
           if no results present for a given unit
         """
         try:
             return self._unit_result_mapping[protocol_unit]
         except KeyError:
-            raise KeyError("No such `protocol_unit` present")
+            raise MissingProtocolUnitError(f"No such `protocol_unit`:{protocol_unit} present")
 
     def result_to_unit(self, protocol_unit_result: ProtocolUnitResult) -> ProtocolUnit:
         try:
             return self._result_unit_mapping[protocol_unit_result]
         except KeyError:
-            raise KeyError("No such `protocol_unit_result` present")
+            raise MissingProtocolUnitError(f"No such `protocol_unit_result`:{protocol_unit_result} present")
 
     def ok(self) -> bool:
         # ensure that for every protocol unit, there is an OK result object

gufe/tests/test_protocol.py

@@ -24,6 +24,8 @@
     ProtocolUnit,
     ProtocolUnitFailure,
     ProtocolUnitResult,
+    MissingProtocolUnitError,
+    ProtocolUnitFailureError
 )
 from gufe.protocols.protocoldag import execute_DAG
 
@@ -692,7 +694,7 @@ def test_missing_result(self, units, successes, failures):
 
         assert not dagresult.ok()
 
-        with pytest.raises(KeyError, match="No success for `protocol_unit` found") as e:
+        with pytest.raises(ProtocolUnitFailureError, match="No success for `protocol_unit`:NoDepUnit\(None\) found"):
             dagresult.unit_to_result(units[2])
 
     def test_plenty_of_fails(self, units, successes, failures):
@@ -721,11 +723,11 @@ def test_foreign_objects(self, units, successes):
             transformation_key=None,
         )
 
-        with pytest.raises(KeyError, match="No such `protocol_unit` present"):
+        with pytest.raises(MissingProtocolUnitError, match="No such `protocol_unit`:NoDepUnit\(None\) present"):
             dagresult.unit_to_result(units[2])
-        with pytest.raises(KeyError, match="No such `protocol_unit` present"):
+        with pytest.raises(MissingProtocolUnitError, match="No such `protocol_unit`:NoDepUnit\(None\) present"):
             dagresult.unit_to_all_results(units[2])
-        with pytest.raises(KeyError, match="No such `protocol_unit_result` present"):
+        with pytest.raises(MissingProtocolUnitError, match="No such `protocol_unit_result`:ProtocolUnitResult\(None\) present"):
             dagresult.result_to_unit(successes[2])