Clean up some types

Author: stan-janssen

shapeshifter_uftp/client/agr_cro_client.py

@@ -1,4 +1,9 @@
-from ..uftp import AgrPortfolioQuery, AgrPortfolioUpdate, PayloadMessageResponse
+from ..uftp import (
+    AgrPortfolioQuery,
+    AgrPortfolioUpdate,
+    PayloadMessageResponse,
+    UsefRole,
+)
 from .base_client import ShapeshifterClient
 
 
@@ -7,8 +12,8 @@ class ShapeshifterAgrCroClient(ShapeshifterClient):
     Client that allows the Aggregator to connect to the CRO.
     """
 
-    sender_role = "AGR"
-    recipient_role = "CRO"
+    sender_role = UsefRole.AGR
+    recipient_role = UsefRole.CRO
 
     def send_agr_portfolio_update(self, message: AgrPortfolioUpdate) -> PayloadMessageResponse:
         """

shapeshifter_uftp/client/agr_dso_client.py

@@ -8,6 +8,7 @@
     FlexSettlementResponse,
     Metering,
     PayloadMessageResponse,
+    UsefRole,
 )
 from .base_client import ShapeshifterClient
 
@@ -17,8 +18,8 @@ class ShapeshifterAgrDsoClient(ShapeshifterClient):
     Client that allows the Aggregator to connect to the DSO.
     """
 
-    sender_role = "AGR"
-    recipient_role = "DSO"
+    sender_role = UsefRole.AGR
+    recipient_role = UsefRole.DSO
 
     def send_d_prognosis(self, message: DPrognosis) -> PayloadMessageResponse:
         """

shapeshifter_uftp/client/base_client.py

@@ -11,16 +11,16 @@
 from ..exceptions import ClientTransportException
 from ..logging import logger
 from ..oauth import OAuthClient, PassthroughOAuthClient
-from ..uftp import PayloadMessage, PayloadMessageResponse, SignedMessage
+from ..uftp import PayloadMessage, PayloadMessageResponse, SignedMessage, UsefRole
 
 
 class ShapeshifterClient:
     """
     Basis for all Shapeshifter client.
     """
 
-    sender_role: str
-    recipient_role: str
+    sender_role: UsefRole
+    recipient_role: UsefRole
     num_outgoing_workers = 10
     num_delivery_attempts = 10
     request_timeout = 30
@@ -32,9 +32,9 @@ def __init__(
         sender_domain: str,
         signing_key: str,
         recipient_domain: str,
-        recipient_endpoint: str = None,
-        recipient_signing_key: str = None,
-        oauth_client: OAuthClient = None,
+        recipient_endpoint: str | None = None,
+        recipient_signing_key: str | None = None,
+        oauth_client: OAuthClient | None = None,
         version: str = "3.1.0"
     ):
         """
@@ -101,7 +101,6 @@ def _send_message(self, message: PayloadMessage) -> PayloadMessageResponse:
         # properties can be calculated in the framework anyway.
         message.version = self.version
         message.sender_domain = self.sender_domain
-        message.sender_role = self.sender_role
         message.recipient_domain = self.recipient_domain
         message.time_stamp = (
             message.time_stamp or datetime.now(timezone.utc).isoformat()

shapeshifter_uftp/client/cro_agr_client.py

@@ -1,4 +1,4 @@
-from ..uftp import AgrPortfolioQueryResponse, AgrPortfolioUpdateResponse
+from ..uftp import AgrPortfolioQueryResponse, AgrPortfolioUpdateResponse, UsefRole
 from .base_client import ShapeshifterClient
 
 
@@ -7,8 +7,8 @@ class ShapeshifterCroAgrClient(ShapeshifterClient):
     Client that allows the CRO to connect to the Aggregator.
     """
 
-    sender_role = "CRO"
-    recipient_role = "AGR"
+    sender_role = UsefRole.CRO
+    recipient_role = UsefRole.AGR
 
     def send_agr_portfolio_update_response(self, message: AgrPortfolioUpdateResponse):
         """

shapeshifter_uftp/client/cro_dso_client.py

@@ -1,4 +1,4 @@
-from ..uftp import DsoPortfolioQueryResponse, DsoPortfolioUpdateResponse
+from ..uftp import DsoPortfolioQueryResponse, DsoPortfolioUpdateResponse, UsefRole
 from .base_client import ShapeshifterClient
 
 
@@ -15,8 +15,8 @@ class ShapeshifterCroDsoClient(ShapeshifterClient):
     DsoPortfolioQuery, respectively.
     """
 
-    sender_role = "CRO"
-    recipient_role = "DSO"
+    sender_role = UsefRole.CRO
+    recipient_role = UsefRole.DSO
 
     def send_dso_portfolio_update_response(self, message: DsoPortfolioUpdateResponse):
         """

shapeshifter_uftp/client/dso_agr_client.py

@@ -8,6 +8,7 @@
     FlexSettlement,
     MeteringResponse,
     PayloadMessageResponse,
+    UsefRole,
 )
 from .base_client import ShapeshifterClient
 
@@ -17,8 +18,8 @@ class ShapeshifterDsoAgrClient(ShapeshifterClient):
     Client that allows the DSO to connect to the Aggregator.
     """
 
-    sender_role = "DSO"
-    recipient_role = "AGR"
+    sender_role = UsefRole.DSO
+    recipient_role = UsefRole.AGR
 
     def send_d_prognosis_response(self, message: DPrognosisResponse) -> PayloadMessageResponse:
         """

shapeshifter_uftp/client/dso_cro_client.py

@@ -1,4 +1,4 @@
-from ..uftp import DsoPortfolioQuery, DsoPortfolioUpdate
+from ..uftp import DsoPortfolioQuery, DsoPortfolioUpdate, UsefRole
 from .base_client import ShapeshifterClient
 
 
@@ -7,8 +7,8 @@ class ShapeshifterDsoCroClient(ShapeshifterClient):
     Client that allows the DSO to connect to the CRO.
     """
 
-    sender_role = "DSO"
-    recipient_role = "CRO"
+    sender_role = UsefRole.DSO
+    recipient_role = UsefRole.CRO
 
     def send_dso_portfolio_update(self, message: DsoPortfolioUpdate):
         """

shapeshifter_uftp/service/base_service.py

@@ -1,9 +1,7 @@
 import re
 from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime, timezone
 from threading import Thread
 from time import sleep
-from uuid import uuid4
 
 import uvicorn
 from fastapi import FastAPI, Response
@@ -22,7 +20,6 @@
 from ..uftp import (
     AcceptedRejected,
     PayloadMessage,
-    PayloadMessageResponse,
     SignedMessage,
     request_response_map,
 )
@@ -154,7 +151,7 @@ def stop(self):
     #              Shapeshifter UFTP implementation.               #
     # ------------------------------------------------------------ #
 
-    def _receive_message(self, message: SignedMessage) -> None:
+    def _receive_message(self, message: SignedMessage) -> Response:
         """
         The default entrypoint for the route. This will unpack the
         message and validate the signature. It will thes pass the

shapeshifter_uftp/transport.py

@@ -22,7 +22,7 @@
     ServiceDiscoveryException,
 )
 from .logging import logger
-from .uftp import PayloadMessage
+from .uftp import PayloadMessage, SignedMessage
 
 _context = XmlContext()
 serializer = XmlSerializer(context=_context, config=SerializerConfig(indent="  "))
@@ -76,7 +76,7 @@ def unseal_message(message: bytes, public_key: str) -> PayloadMessage:
         raise SchemaException(str(exc)) from exc
 
 
-def to_xml(message: PayloadMessage) -> str:
+def to_xml(message: PayloadMessage | SignedMessage) -> str:
     """
     Serialize the given PayloadMessage into an XML string.
     """
@@ -174,7 +174,7 @@ def get_keys(domain, role):
         )
 
     # Verify that the string is of the expected length (4 + 44 bytes or 4 + 88 bytes)
-    if not len(result) in (48, 92):
+    if len(result) not in (48, 92):
         raise AuthenticationTimeoutException(
             f"Could not retrieve public key(s) at {dns_name}: "
             f"string '{result}' was not of appropriate length (48 or 90 characters)"
@@ -190,7 +190,7 @@ def get_keys(domain, role):
         ) from exc
 
     # Now verify that the decoded length is 64
-    if not len(combined_keys) in (32, 64):
+    if len(combined_keys) not in (32, 64):
         raise AuthenticationTimeoutException(
             f"Could not retrieve public keys at {dns_name}: "
             f"decoded base64 data should be 32 or 64 bytes long, "