fix(parser): overload parse when using envelope (#885)

Author: heitorlessa

Diff for: aws_lambda_powertools/utilities/parser/envelopes/base.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, Type, TypeVar, Union
 
-from ..types import Model
+from aws_lambda_powertools.utilities.parser.types import Model
 
 logger = logging.getLogger(__name__)
 

Diff for: aws_lambda_powertools/utilities/parser/parser.py

@@ -1,25 +1,23 @@
 import logging
-from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union
+from typing import Any, Callable, Dict, Optional, Type, overload
+
+from aws_lambda_powertools.utilities.parser.types import EnvelopeModel, EventParserReturnType, Model
 
 from ...middleware_factory import lambda_handler_decorator
 from ..typing import LambdaContext
 from .envelopes.base import Envelope
 from .exceptions import InvalidEnvelopeError, InvalidModelTypeError
-from .types import Model
 
 logger = logging.getLogger(__name__)
 
 
-EventParserReturnType = TypeVar("EventParserReturnType")
-
-
 @lambda_handler_decorator
 def event_parser(
     handler: Callable[[Any, LambdaContext], EventParserReturnType],
     event: Dict[str, Any],
     context: LambdaContext,
     model: Type[Model],
-    envelope: Optional[Union[Envelope, Type[Envelope]]] = None,
+    envelope: Optional[Type[Envelope]] = None,
 ) -> EventParserReturnType:
     """Lambda handler decorator to parse & validate events using Pydantic models
 
@@ -81,14 +79,22 @@ def handler(event: Order, context: LambdaContext):
     InvalidEnvelopeError
         When envelope given does not implement BaseEnvelope
     """
-    parsed_event = parse(event=event, model=model, envelope=envelope)
+    parsed_event = parse(event=event, model=model, envelope=envelope) if envelope else parse(event=event, model=model)
     logger.debug(f"Calling handler {handler.__name__}")
     return handler(parsed_event, context)
 
 
-def parse(
-    event: Dict[str, Any], model: Type[Model], envelope: Optional[Union[Envelope, Type[Envelope]]] = None
-) -> Model:
+@overload
+def parse(event: Dict[str, Any], model: Type[Model]) -> Model:
+    ...
+
+
+@overload
+def parse(event: Dict[str, Any], model: Type[Model], envelope: Type[Envelope]) -> EnvelopeModel:
+    ...
+
+
+def parse(event: Dict[str, Any], model: Type[Model], envelope: Optional[Type[Envelope]] = None):
     """Standalone function to parse & validate events using Pydantic models
 
     Typically used when you need fine-grained control over error handling compared to event_parser decorator.

Diff for: aws_lambda_powertools/utilities/parser/types.py

@@ -12,3 +12,5 @@
     from typing_extensions import Literal  # noqa: F401
 
 Model = TypeVar("Model", bound=BaseModel)
+EnvelopeModel = TypeVar("EnvelopeModel")
+EventParserReturnType = TypeVar("EventParserReturnType")