Leverage str, enum capability of Python for any generated enums (#133)

* str_enum: Change generation of S2 specification to include the subclass of an enum.

* str_enum: Exclude generated files from pyright check.

* str_enum: ombc classes should be explicitely exported in init.

* str_enum: add missing Any to dict type in S2Parser.

Author: lfse-slafleur

ci/generate_s2.sh

@@ -2,4 +2,4 @@
 
 
 . .venv/bin/activate
-datamodel-codegen  --input specification/openapi.yml --input-file-type openapi --output-model-type pydantic_v2.BaseModel --output src/s2python/generated/gen_s2.py --use-one-literal-as-default
+datamodel-codegen  --input specification/openapi.yml --input-file-type openapi --output-model-type pydantic_v2.BaseModel --output src/s2python/generated/gen_s2.py --use-one-literal-as-default --use-subclass-enum --output-datetime-class AwareDatetime --use-double-quotes

pyproject.toml

@@ -64,3 +64,6 @@ docs = [
 
 [project.scripts]
 s2python = "s2python.tools.cli:s2python_cmd"
+
+[tool.black]
+line-length=100

pyrightconfig.json

@@ -4,6 +4,10 @@
     "tests"
   ],
 
+  "exclude": [
+    "src/s2python/generated/"
+  ],
+
   "defineConstant": {
     "DEBUG": true
   }

src/s2python/generated/gen_s2.py

@@ -1,34 +1,21 @@
 # generated by datamodel-codegen:
 #   filename:  openapi.yml
-#   timestamp: 2024-07-29T10:18:52+00:00
+#   timestamp: 2025-07-29T10:12:16+00:00
 
 from __future__ import annotations
 
 from enum import Enum
-from typing import List, Optional
+from typing import List, Literal, Optional
 
-from pydantic import (
-    AwareDatetime,
-    BaseModel,
-    ConfigDict,
-    Field,
-    RootModel,
-    conint,
-    constr,
-)
-from typing_extensions import Literal
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, RootModel, conint, constr
 
 
 class Duration(RootModel[conint(ge=0)]):
-    root: conint(ge=0) = Field(  # pyright: ignore[reportInvalidTypeForm]
-        ..., description="Duration in milliseconds"
-    )
+    root: conint(ge=0) = Field(..., description="Duration in milliseconds")
 
 
 class ID(RootModel[constr(pattern=r"[a-zA-Z0-9\-_:]{2,64}")]):
-    root: constr(pattern=r"[a-zA-Z0-9\-_:]{2,64}") = (  # pyright: ignore[reportInvalidTypeForm]
-        Field(..., description="UUID")
-    )
+    root: constr(pattern=r"[a-zA-Z0-9\-_:]{2,64}") = Field(..., description="UUID")
 
 
 class Currency(Enum):
@@ -134,12 +121,12 @@ class Currency(Enum):
     ZWL = "ZWL"
 
 
-class SessionRequestType(Enum):
+class SessionRequestType(str, Enum):
     RECONNECT = "RECONNECT"
     TERMINATE = "TERMINATE"
 
 
-class RevokableObjects(Enum):
+class RevokableObjects(str, Enum):
     PEBC_PowerConstraints = "PEBC.PowerConstraints"
     PEBC_EnergyConstraint = "PEBC.EnergyConstraint"
     PEBC_Instruction = "PEBC.Instruction"
@@ -155,12 +142,12 @@ class RevokableObjects(Enum):
     DDBC_Instruction = "DDBC.Instruction"
 
 
-class EnergyManagementRole(Enum):
+class EnergyManagementRole(str, Enum):
     CEM = "CEM"
     RM = "RM"
 
 
-class ReceptionStatusValues(Enum):
+class ReceptionStatusValues(str, Enum):
     INVALID_DATA = "INVALID_DATA"
     INVALID_MESSAGE = "INVALID_MESSAGE"
     INVALID_CONTENT = "INVALID_CONTENT"
@@ -234,8 +221,7 @@ class Timer(BaseModel):
         description="Human readable name/description of the Timer. This element is only intended for diagnostic purposes and not for HMI applications.",
     )
     duration: Duration = Field(
-        ...,
-        description="The time it takes for the Timer to finish after it has been started",
+        ..., description="The time it takes for the Timer to finish after it has been started"
     )
 
 
@@ -381,20 +367,20 @@ class DDBCAverageDemandRateForecastElement(BaseModel):
     )
 
 
-class RoleType(Enum):
+class RoleType(str, Enum):
     ENERGY_PRODUCER = "ENERGY_PRODUCER"
     ENERGY_CONSUMER = "ENERGY_CONSUMER"
     ENERGY_STORAGE = "ENERGY_STORAGE"
 
 
-class Commodity(Enum):
+class Commodity(str, Enum):
     GAS = "GAS"
     HEAT = "HEAT"
     ELECTRICITY = "ELECTRICITY"
     OIL = "OIL"
 
 
-class CommodityQuantity(Enum):
+class CommodityQuantity(str, Enum):
     ELECTRIC_POWER_L1 = "ELECTRIC.POWER.L1"
     ELECTRIC_POWER_L2 = "ELECTRIC.POWER.L2"
     ELECTRIC_POWER_L3 = "ELECTRIC.POWER.L3"
@@ -407,7 +393,7 @@ class CommodityQuantity(Enum):
     OIL_FLOW_RATE = "OIL.FLOW_RATE"
 
 
-class InstructionStatus(Enum):
+class InstructionStatus(str, Enum):
     NEW = "NEW"
     ACCEPTED = "ACCEPTED"
     REJECTED = "REJECTED"
@@ -417,7 +403,7 @@ class InstructionStatus(Enum):
     ABORTED = "ABORTED"
 
 
-class ControlType(Enum):
+class ControlType(str, Enum):
     POWER_ENVELOPE_BASED_CONTROL = "POWER_ENVELOPE_BASED_CONTROL"
     POWER_PROFILE_BASED_CONTROL = "POWER_PROFILE_BASED_CONTROL"
     OPERATION_MODE_BASED_CONTROL = "OPERATION_MODE_BASED_CONTROL"
@@ -427,17 +413,17 @@ class ControlType(Enum):
     NO_SELECTION = "NO_SELECTION"
 
 
-class PEBCPowerEnvelopeLimitType(Enum):
+class PEBCPowerEnvelopeLimitType(str, Enum):
     UPPER_LIMIT = "UPPER_LIMIT"
     LOWER_LIMIT = "LOWER_LIMIT"
 
 
-class PEBCPowerEnvelopeConsequenceType(Enum):
+class PEBCPowerEnvelopeConsequenceType(str, Enum):
     VANISH = "VANISH"
     DEFER = "DEFER"
 
 
-class PPBCPowerSequenceStatus(Enum):
+class PPBCPowerSequenceStatus(str, Enum):
     NOT_SCHEDULED = "NOT_SCHEDULED"
     SCHEDULED = "SCHEDULED"
     EXECUTING = "EXECUTING"
@@ -507,8 +493,7 @@ class SessionRequest(BaseModel):
     message_id: ID
     request: SessionRequestType = Field(..., description="The type of request")
     diagnostic_label: Optional[str] = Field(
-        None,
-        description="Optional field for a human readible descirption for debugging purposes",
+        None, description="Optional field for a human readible descirption for debugging purposes"
     )
 
 
@@ -586,12 +571,10 @@ class PEBCEnergyConstraint(BaseModel):
         description="Identifier of this PEBC.EnergyConstraints. Must be unique in the scope of the Resource Manager, for at least the duration of the session between Resource Manager and CEM.",
     )
     valid_from: AwareDatetime = Field(
-        ...,
-        description="Moment this PEBC.EnergyConstraints information starts to be valid",
+        ..., description="Moment this PEBC.EnergyConstraints information starts to be valid"
     )
     valid_until: AwareDatetime = Field(
-        ...,
-        description="Moment until this PEBC.EnergyConstraints information is valid.",
+        ..., description="Moment until this PEBC.EnergyConstraints information is valid."
     )
     upper_average_power: float = Field(
         ...,
@@ -634,8 +617,7 @@ class PPBCScheduleInstruction(BaseModel):
         description="Indicates the moment the PPBC.PowerSequence shall start. When the specified execution time is in the past, execution must start as soon as possible.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition",
+        ..., description="Indicates if this is an instruction during an abnormal condition"
     )
 
 
@@ -665,8 +647,7 @@ class PPBCStartInterruptionInstruction(BaseModel):
         description="Indicates the moment the PPBC.PowerSequence shall be interrupted. When the specified execution time is in the past, execution must start as soon as possible.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition",
+        ..., description="Indicates if this is an instruction during an abnormal condition"
     )
 
 
@@ -697,8 +678,7 @@ class PPBCEndInterruptionInstruction(BaseModel):
         description="Indicates the moment PPBC.PowerSequence interruption shall end. When the specified execution time is in the past, execution must start as soon as possible.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition",
+        ..., description="Indicates if this is an instruction during an abnormal condition"
     )
 
 
@@ -745,8 +725,7 @@ class OMBCInstruction(BaseModel):
         description="The number indicates the factor with which the OMBC.OperationMode should be configured. The factor should be greater than or equal than 0 and less or equal to 1.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition",
+        ..., description="Indicates if this is an instruction during an abnormal condition"
     )
 
 
@@ -824,8 +803,7 @@ class FRBCInstruction(BaseModel):
         description="Indicates the moment the execution of the instruction shall start. When the specified execution time is in the past, execution must start as soon as possible.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition.",
+        ..., description="Indicates if this is an instruction during an abnormal condition."
     )
 
 
@@ -871,8 +849,7 @@ class DDBCActuatorStatus(BaseModel):
     message_id: ID
     actuator_id: ID = Field(..., description="ID of the actuator this messages refers to")
     active_operation_mode_id: ID = Field(
-        ...,
-        description="The operation mode that is presently active for this actuator.",
+        ..., description="The operation mode that is presently active for this actuator."
     )
     operation_mode_factor: float = Field(
         ...,
@@ -903,8 +880,7 @@ class DDBCInstruction(BaseModel):
         description="Indicates the moment the execution of the instruction shall start. When the specified execution time is in the past, execution must start as soon as possible.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition",
+        ..., description="Indicates if this is an instruction during an abnormal condition"
     )
     actuator_id: ID = Field(..., description="ID of the actuator this Instruction belongs to.")
     operation_mode_id: ID = Field(..., description="ID of the DDBC.OperationMode")
@@ -937,8 +913,7 @@ class PowerValue(BaseModel):
         ..., description="The power quantity the value refers to"
     )
     value: float = Field(
-        ...,
-        description="Power value expressed in the unit associated with the CommodityQuantity",
+        ..., description="Power value expressed in the unit associated with the CommodityQuantity"
     )
 
 
@@ -1020,8 +995,7 @@ class PEBCAllowedLimitRange(BaseModel):
         ..., description="Type of power quantity this PEBC.AllowedLimitRange applies to"
     )
     limit_type: PEBCPowerEnvelopeLimitType = Field(
-        ...,
-        description="Indicates if this ranges applies to the upper limit or the lower limit",
+        ..., description="Indicates if this ranges applies to the upper limit or the lower limit"
     )
     range_boundary: NumberRange = Field(
         ...,
@@ -1193,8 +1167,7 @@ class ResourceManagerDetails(BaseModel):
     )
     manufacturer: Optional[str] = Field(None, description="Name of Manufacturer")
     model: Optional[str] = Field(
-        None,
-        description="Name of the model of the device (provided by the manufacturer)",
+        None, description="Name of the model of the device (provided by the manufacturer)"
     )
     serial_number: Optional[str] = Field(
         None, description="Serial number of the device (provided by the manufacturer)"
@@ -1218,8 +1191,7 @@ class ResourceManagerDetails(BaseModel):
         description="Currency to be used for all information regarding costs. Mandatory if cost information is published.",
     )
     provides_forecast: bool = Field(
-        ...,
-        description="Indicates whether the ResourceManager is able to provide PowerForecasts",
+        ..., description="Indicates whether the ResourceManager is able to provide PowerForecasts"
     )
     provides_power_measurement_types: List[CommodityQuantity] = Field(
         ...,
@@ -1306,8 +1278,7 @@ class PEBCInstruction(BaseModel):
         description="Indicates the moment the execution of the instruction shall start. When the specified execution time is in the past, execution must start as soon as possible.",
     )
     abnormal_condition: bool = Field(
-        ...,
-        description="Indicates if this is an instruction during an abnormal condition.",
+        ..., description="Indicates if this is an instruction during an abnormal condition."
     )
     power_constraints_id: ID = Field(
         ...,
@@ -1380,8 +1351,7 @@ class PPBCPowerSequence(BaseModel):
         min_length=1,
     )
     is_interruptible: bool = Field(
-        ...,
-        description="Indicates whether the option of pausing a sequence is available.",
+        ..., description="Indicates whether the option of pausing a sequence is available."
     )
     max_pause_before: Optional[Duration] = Field(
         None,
@@ -1509,10 +1479,7 @@ class FRBCActuatorDescription(BaseModel):
         description="Human readable name/description for the actuator. This element is only intended for diagnostic purposes and not for HMI applications.",
     )
     supported_commodities: List[Commodity] = Field(
-        ...,
-        description="List of all supported Commodities.",
-        max_length=4,
-        min_length=1,
+        ..., description="List of all supported Commodities.", max_length=4, min_length=1
     )
     operation_modes: List[FRBCOperationMode] = Field(
         ...,

src/s2python/ombc/__init__.py

@@ -3,3 +3,11 @@
 from s2python.ombc.ombc_status import OMBCStatus
 from s2python.ombc.ombc_system_description import OMBCSystemDescription
 from s2python.ombc.ombc_timer_status import OMBCTimerStatus
+
+__all__ = [
+    "OMBCInstruction",
+    "OMBCOperationMode",
+    "OMBCStatus",
+    "OMBCSystemDescription",
+    "OMBCTimerStatus",
+]

src/s2python/s2_parser.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import Optional, TypeVar, Union, Type, Dict
+from typing import Optional, TypeVar, Union, Type, Dict, Any
 
 from s2python.common import (
     Handshake,
@@ -67,13 +67,13 @@
 
 class S2Parser:
     @staticmethod
-    def _parse_json_if_required(unparsed_message: Union[dict, str, bytes]) -> dict:
+    def _parse_json_if_required(unparsed_message: Union[dict[Any, Any], str, bytes]) -> dict:
         if isinstance(unparsed_message, (str, bytes)):
             return json.loads(unparsed_message)
         return unparsed_message
 
     @staticmethod
-    def parse_as_any_message(unparsed_message: Union[dict, str, bytes]) -> S2Message:
+    def parse_as_any_message(unparsed_message: Union[dict[Any, Any], str, bytes]) -> S2Message:
         """Parse the message as any S2 python message regardless of message type.
 
         :param unparsed_message: The message as a JSON-formatted string or as a json-parsed dictionary.
@@ -94,7 +94,7 @@ def parse_as_any_message(unparsed_message: Union[dict, str, bytes]) -> S2Message
 
     @staticmethod
     def parse_as_message(
-        unparsed_message: Union[dict, str, bytes], as_message: Type[M]
+        unparsed_message: Union[dict[Any, Any], str, bytes], as_message: Type[M]
     ) -> M:
         """Parse the message to a specific S2 python message.
 
@@ -108,7 +108,7 @@ def parse_as_message(
 
     @staticmethod
     def parse_message_type(
-        unparsed_message: Union[dict, str, bytes],
+        unparsed_message: Union[dict[Any, Any], str, bytes],
     ) -> Optional[S2MessageType]:
         """Parse only the message type from the unparsed message.