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

Author: heitorlessa

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__)
 

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.

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")