fix(parser): mypy support for payload type override as models (#883)

Co-authored-by: Tom McCarthy <[email protected]>

aws_lambda_powertools/utilities/parser/models/alb.py

@@ -1,7 +1,9 @@
-from typing import Dict
+from typing import Dict, Union
 
 from pydantic import BaseModel
 
+from aws_lambda_powertools.utilities.parser.types import Model
+
 
 class AlbRequestContextData(BaseModel):
     targetGroupArn: str
@@ -14,7 +16,7 @@ class AlbRequestContext(BaseModel):
 class AlbModel(BaseModel):
     httpMethod: str
     path: str
-    body: str
+    body: Union[str, Model]
     isBase64Encoded: bool
     headers: Dict[str, str]
     queryStringParameters: Dict[str, str]

aws_lambda_powertools/utilities/parser/models/apigw.py

@@ -1,10 +1,10 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, root_validator
 from pydantic.networks import IPvAnyNetwork
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal, Model
 
 
 class ApiGatewayUserCertValidity(BaseModel):
@@ -89,4 +89,4 @@ class APIGatewayProxyEventModel(BaseModel):
     pathParameters: Optional[Dict[str, str]]
     stageVariables: Optional[Dict[str, str]]
     isBase64Encoded: bool
-    body: Optional[str]
+    body: Optional[Union[str, Model]]

aws_lambda_powertools/utilities/parser/models/apigwv2.py

@@ -1,10 +1,10 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field
 from pydantic.networks import IPvAnyNetwork
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal, Model
 
 
 class RequestContextV2AuthorizerIamCognito(BaseModel):
@@ -67,5 +67,5 @@ class APIGatewayProxyEventV2Model(BaseModel):
     pathParameters: Optional[Dict[str, str]]
     stageVariables: Optional[Dict[str, str]]
     requestContext: RequestContextV2
-    body: Optional[str]
+    body: Optional[Union[str, Model]]
     isBase64Encoded: bool

aws_lambda_powertools/utilities/parser/models/cloudwatch.py

@@ -3,17 +3,19 @@
 import logging
 import zlib
 from datetime import datetime
-from typing import List
+from typing import List, Union
 
 from pydantic import BaseModel, Field, validator
 
+from aws_lambda_powertools.utilities.parser.types import Model
+
 logger = logging.getLogger(__name__)
 
 
 class CloudWatchLogsLogEvent(BaseModel):
     id: str  # noqa AA03 VNE003
     timestamp: datetime
-    message: str
+    message: Union[str, Model]
 
 
 class CloudWatchLogsDecode(BaseModel):

aws_lambda_powertools/utilities/parser/models/dynamodb.py

@@ -1,16 +1,16 @@
 from datetime import date
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal, Model
 
 
 class DynamoDBStreamChangedRecordModel(BaseModel):
     ApproximateCreationDateTime: Optional[date]
     Keys: Dict[str, Dict[str, Any]]
-    NewImage: Optional[Dict[str, Any]]
-    OldImage: Optional[Dict[str, Any]]
+    NewImage: Optional[Union[Dict[str, Any], Model]]
+    OldImage: Optional[Union[Dict[str, Any], Model]]
     SequenceNumber: str
     SizeBytes: int
     StreamViewType: Literal["NEW_AND_OLD_IMAGES", "KEYS_ONLY", "NEW_IMAGE", "OLD_IMAGE"]

aws_lambda_powertools/utilities/parser/models/event_bridge.py

@@ -1,8 +1,10 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field
 
+from aws_lambda_powertools.utilities.parser.types import Model
+
 
 class EventBridgeModel(BaseModel):
     version: str
@@ -13,5 +15,5 @@ class EventBridgeModel(BaseModel):
     region: str
     resources: List[str]
     detail_type: str = Field(None, alias="detail-type")
-    detail: Dict[str, Any]
+    detail: Union[Dict[str, Any], Model]
     replay_name: Optional[str] = Field(None, alias="replay-name")

aws_lambda_powertools/utilities/parser/models/kinesis.py

@@ -1,12 +1,12 @@
 import base64
 import logging
 from binascii import Error as BinAsciiError
-from typing import List
+from typing import List, Union
 
 from pydantic import BaseModel, validator
 from pydantic.types import PositiveInt
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal, Model
 
 logger = logging.getLogger(__name__)
 
@@ -15,7 +15,7 @@ class KinesisDataStreamRecordPayload(BaseModel):
     kinesisSchemaVersion: str
     partitionKey: str
     sequenceNumber: PositiveInt
-    data: bytes  # base64 encoded str is parsed into bytes
+    data: Union[bytes, Model]  # base64 encoded str is parsed into bytes
     approximateArrivalTimestamp: float
 
     @validator("data", pre=True, allow_reuse=True)

aws_lambda_powertools/utilities/parser/models/s3.py

@@ -6,7 +6,7 @@
 from pydantic.networks import IPvAnyNetwork
 from pydantic.types import NonNegativeFloat
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal
 
 
 class S3EventRecordGlacierRestoreEventData(BaseModel):

aws_lambda_powertools/utilities/parser/models/s3_object_event.py

@@ -1,7 +1,9 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 from pydantic import BaseModel, HttpUrl
 
+from aws_lambda_powertools.utilities.parser.types import Model
+
 
 class S3ObjectContext(BaseModel):
     inputS3Url: HttpUrl
@@ -12,7 +14,7 @@ class S3ObjectContext(BaseModel):
 class S3ObjectConfiguration(BaseModel):
     accessPointArn: str
     supportingAccessPointArn: str
-    payload: str
+    payload: Union[str, Model]
 
 
 class S3ObjectUserRequest(BaseModel):

aws_lambda_powertools/utilities/parser/models/sns.py

@@ -1,10 +1,10 @@
 from datetime import datetime
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, root_validator
 from pydantic.networks import HttpUrl
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal, Model
 
 
 class SnsMsgAttributeModel(BaseModel):
@@ -18,7 +18,7 @@ class SnsNotificationModel(BaseModel):
     UnsubscribeUrl: HttpUrl
     Type: Literal["Notification"]
     MessageAttributes: Optional[Dict[str, SnsMsgAttributeModel]]
-    Message: str
+    Message: Union[str, Model]
     MessageId: str
     SigningCertUrl: HttpUrl
     Signature: str

aws_lambda_powertools/utilities/parser/models/sqs.py

@@ -1,9 +1,9 @@
 from datetime import datetime
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel
 
-from ..types import Literal
+from aws_lambda_powertools.utilities.parser.types import Literal, Model
 
 
 class SqsAttributesModel(BaseModel):
@@ -52,7 +52,7 @@ class SqsMsgAttributeModel(BaseModel):
 class SqsRecordModel(BaseModel):
     messageId: str
     receiptHandle: str
-    body: str
+    body: Union[str, Model]
     attributes: SqsAttributesModel
     messageAttributes: Dict[str, SqsMsgAttributeModel]
     md5OfBody: str

docs/utilities/parser.md

@@ -170,6 +170,8 @@ Parser comes with the following built-in models:
 
 You can extend them to include your own models, and yet have all other known fields parsed along the way.
 
+!!! tip "For Mypy users, we only allow type override for fields where payload is injected e.g. `detail`, `body`, etc."
+
 **EventBridge example**
 
 === "extending_builtin_models.py"