diff --git a/src/s2python/s2_parser.py b/src/s2python/s2_parser.py index e5b7ea1..c5f6c14 100644 --- a/src/s2python/s2_parser.py +++ b/src/s2python/s2_parser.py @@ -24,6 +24,7 @@ FRBCTimerStatus, FRBCUsageForecast, ) +from s2python.pebc import PEBCPowerConstraints, PEBCEnergyConstraint, PEBCInstruction from s2python.ppbc import PPBCScheduleInstruction from s2python.message import S2Message @@ -48,6 +49,9 @@ "FRBC.TimerStatus": FRBCTimerStatus, "FRBC.UsageForecast": FRBCUsageForecast, "PPBC.ScheduleInstruction": PPBCScheduleInstruction, + "PEBC.PowerConstraints": PEBCPowerConstraints, + "PEBC.Instruction": PEBCInstruction, + "PEBC.EnergyConstraint": PEBCEnergyConstraint, "Handshake": Handshake, "HandshakeResponse": HandshakeResponse, "InstructionStatusUpdate": InstructionStatusUpdate, @@ -90,7 +94,9 @@ def parse_as_any_message(unparsed_message: Union[dict, str, bytes]) -> S2Message return TYPE_TO_MESSAGE_CLASS[message_type].model_validate(message_json) @staticmethod - def parse_as_message(unparsed_message: Union[dict, str, bytes], as_message: Type[M]) -> M: + def parse_as_message( + unparsed_message: Union[dict, str, bytes], as_message: Type[M] + ) -> M: """Parse the message to a specific S2 python message. :param unparsed_message: The message as a JSON-formatted string or as a JSON-parsed dictionary.