fix(middleware_factory): ret type annotation for handler dec

[huonw]

Issue #, if available: #1060

Description of changes:

This is two versions of typing lambda_handler_decorator more accurately for #1060, to ensure that middleware using it have a type. They end up with type Callable[..., Any] but that's better than nothing.

The first commit 937e6fe tries to be more accurate via overloads and protocols, but it doesn't work as written, and I cannot quite work out why. The errors below demonstrate the issues. It's potentially resolvable with 3.10's ParamSpec (I'm not sure typing_extensions depended upon for 3.8 and 3.8?):

Errors

aws_lambda_powertools/utilities/validation/validator.py:12:2: error: No overload variant of "lambda_handler_decorator" matches argument type "Callable[[Callable[..., Any], Union[Dict[Any, Any], str], Any, Optional[Dict[Any, Any]], Optional[Dict[Any, Any]], Optional[Dict[Any, Any]], Optional[Dict[Any, Any]], str, Optional[Dict[Any, Any]]], Any]"  [call-overload]
aws_lambda_powertools/utilities/validation/validator.py:12:2: note: Possible overload variants:
aws_lambda_powertools/utilities/validation/validator.py:12:2: note:     def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator
aws_lambda_powertools/utilities/validation/validator.py:12:2: note:     def lambda_handler_decorator(decorator: None = ..., trace_execution: Optional[bool] = ...) -> Callable[[_FactoryDecorator], _HandlerDecorator]
aws_lambda_powertools/utilities/parser/parser.py:14:2: error: No overload variant of "lambda_handler_decorator" matches argument type "Callable[[Callable[[Any, LambdaContext], EventParserReturnType], Dict[str, Any], LambdaContext, Type[Model], Optional[Type[Envelope]]], EventParserReturnType]"  [call-overload]
aws_lambda_powertools/utilities/parser/parser.py:14:2: note: Possible overload variants:
aws_lambda_powertools/utilities/parser/parser.py:14:2: note:     def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator
aws_lambda_powertools/utilities/parser/parser.py:14:2: note:     def lambda_handler_decorator(decorator: None = ..., trace_execution: Optional[bool] = ...) -> Callable[[_FactoryDecorator], _HandlerDecorator]
aws_lambda_powertools/utilities/idempotency/idempotency.py:20:2: error: No overload variant of "lambda_handler_decorator" matches argument type "Callable[[Callable[[Any, LambdaContext], Any], Dict[str, Any], LambdaContext, BasePersistenceLayer, Optional[IdempotencyConfig], KwArg(Any)], Any]"  [call-overload]
aws_lambda_powertools/utilities/idempotency/idempotency.py:20:2: note: Possible overload variants:
aws_lambda_powertools/utilities/idempotency/idempotency.py:20:2: note:     def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator
aws_lambda_powertools/utilities/idempotency/idempotency.py:20:2: note:     def lambda_handler_decorator(decorator: None = ..., trace_execution: Optional[bool] = ...) -> Callable[[_FactoryDecorator], _HandlerDecorator]
aws_lambda_powertools/utilities/data_classes/event_source.py:8:2: error: No overload variant of "lambda_handler_decorator" matches argument type "Callable[[Callable[[Any, LambdaContext], Any], Dict[str, Any], LambdaContext, Type[DictWrapper]], Any]"  [call-overload]
aws_lambda_powertools/utilities/data_classes/event_source.py:8:2: note: Possible overload variants:
aws_lambda_powertools/utilities/data_classes/event_source.py:8:2: note:     def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator
aws_lambda_powertools/utilities/data_classes/event_source.py:8:2: note:     def lambda_handler_decorator(decorator: None = ..., trace_execution: Optional[bool] = ...) -> Callable[[_FactoryDecorator], _HandlerDecorator]
aws_lambda_powertools/utilities/batch/base.py:156:2: error: No overload variant of "lambda_handler_decorator" matches argument type "Callable[[Callable[..., Any], Dict[Any, Any], Dict[Any, Any], Callable[..., Any], BasePartialProcessor], Any]"  [call-overload]
aws_lambda_powertools/utilities/batch/base.py:156:2: note: Possible overload variants:
aws_lambda_powertools/utilities/batch/base.py:156:2: note:     def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator
aws_lambda_powertools/utilities/batch/base.py:156:2: note:     def lambda_handler_decorator(decorator: None = ..., trace_execution: Optional[bool] = ...) -> Callable[[_FactoryDecorator], _HandlerDecorator]
aws_lambda_powertools/utilities/batch/sqs.py:179:2: error: No overload variant of "lambda_handler_decorator" matches argument type "Callable[[Callable[..., Any], Dict[Any, Any], Dict[Any, Any], Callable[..., Any], Optional[Any], bool, Optional[Any]], Any]"  [call-overload]
aws_lambda_powertools/utilities/batch/sqs.py:179:2: note: Possible overload variants:
aws_lambda_powertools/utilities/batch/sqs.py:179:2: note:     def lambda_handler_decorator(decorator: _FactoryDecorator) -> _HandlerDecorator
aws_lambda_powertools/utilities/batch/sqs.py:179:2: note:     def lambda_handler_decorator(decorator: None = ..., trace_execution: Optional[bool] = ...) -> Callable[[_FactoryDecorator], _HandlerDecorator]

Checklist

• Meet tenets criteria
• Update tests
• Update docs
• PR title follows conventional commit semantics

Breaking change checklist

RFC issue #:

• Migration process documented
• Implement warnings (if it can live side by side)

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[codecov-commenter]

Codecov Report

Merging #1066 (665e0bb) into develop (2a3ff9a) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff            @@
##           develop    #1066   +/-   ##
========================================
  Coverage    99.96%   99.96%           
========================================
  Files          119      119           
  Lines         5365     5365           
  Branches       612      612           
========================================
  Hits          5363     5363           
  Partials         2        2           

Impacted Files	Coverage Δ
...ws_lambda_powertools/middleware_factory/factory.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2a3ff9a...665e0bb. Read the comment docs.

[heitorlessa]

Thanks a lot @huonw !! Question: Do we lose type signature with the current solution? e.g., create a decorator with a given type annotation on a key-value param, call it using an incorrect type and let mypy detect it.

The way I solved this for Tracer (similar, not the same) on preserving the function signature types is using AnyCallableT - This ... works similar to ParamSpec that's available in future version.

Instead of using overload, could you try with AnyCallableT only? I can have a look at this after we fix some bugs in Logger.

[huonw]

I think this is an strict improvement over the current code (i.e. nothing is being lost): the middlewares used to be completely untyped, while now they're at least Callable[..., Any], and thus the handlers they're used on are typed Any.

Are you proposing something like def lambda_handler_decorator(decorator: Optional[AnyCallableT] = None, trace_execution: Optional[bool] = None) -> AnyCallableT. Unfortunately I don't think that'll work quite right, because the type of decorator is different to the return value: decorator has arguments handler, event and context (and kwargs), while the value returned by lambda_handler_decorator just has arguments handler (and kwargs). Let me know if I'm misunderstanding the behaviour here.

[heitorlessa]

If your current Callable type is helping preserve the type signature (not being done today), then we should stop here and merge \o/!

The AnyCallableT suggestion was for the first layer (like Tracer, Logger, Metrics decorators), it wouldn't work for the second one for reasons you called out - hence the thought to experiment as that wasn't the final solution.

Let me know the answer to the types as I'm not entirely sure I understood. If types aren't being preserved with the latest commit, we can experiment other ways (typing from my phone again, need to think of other ways still)

[huonw]

Ah, I see. I think that may be tricky and/or require manual casts at each use site of the decorator:

from typing import Any, Callable, cast

AnyCallable = Callable[..., Any]
def base_decorator(f: AnyCallable) -> AnyCallable:
    return f

@base_decorator
def my_decorator(g: AnyCallable) -> AnyCallable:
    return g

my_decorator = cast(Callable[[AnyCallable], AnyCallable], my_decorator)

There's quite a few downsides of this:

1. requires manually specifying the type at each site
2. using Callable directly like this still doesn't get the type correct right (doesn't distinguish between @decorator like the above and the partially applied version with kwargs @decorator(foo=1) def f() ...: it'd be restricted to only one or the other), although one could probably define an appropriate protocol
3. some tools may complain about having two declarations of my_decorator too (one via the decorator, and one via the cast

On balance, and especially because these decorators are designed for handlers (i.e. functions that typically aren't used elsewhere in library code), I feel the minimal Any-based approach is the best immediate solution.

[heitorlessa]

Looking into this tomorrow ;-) Finished the other bits

[heitorlessa]

Just ran --strict and: Found 842 errors in 73 files (checked 125 source files) - will definitely look into it later. For now, I'll test a custom middleware with an additional kwarg to check whether Mypy can detect it with the new type annotation, then merge ;)

[heitorlessa]

Nah, spent the last two hours and while I made headwinds towards identifying kwargs in a non-strict mode, I hit a dead end with too few arguments later - Adding the patch as ZIP and the sample code I used, in case you know a solution.

Otherwise, I'll take this opportunity to fix a doc typo as part of this PR and then merge.

---

factory.patch.zip

Sample code used

from typing import Any, Callable, Dict, List, Optional

from aws_lambda_powertools.middleware_factory import lambda_handler_decorator


@lambda_handler_decorator
def obfuscate_sensitive_data(
    handler: Callable[..., Any], event: Dict[str, Any], context: Any, fields: Optional[List[str]] = None
) -> Any:
    # Obfuscate email before calling Lambda handler
    if fields:
        for field in fields:
            if field in event:
                event[field] = "REDACTED"

    return handler(event, context)


@obfuscate_sensitive_data(fields=10)
def lambda_handler(event: Dict[str, Any], context: Any) -> Any:
    ...

Mypy strict output against sample code

blah.py:19:2: error: Unexpected keyword argument "fields" for "obfuscate_sensitive_data"  [call-arg]
blah.py:19:2: error: Too few arguments for "obfuscate_sensitive_data"  [call-arg]
Found 2 errors in 1 file (checked 1 source file)

[huonw]

I have no further insight on that patch 😄

[heitorlessa]

suggesting a reword on return type comment, accepting, and merging.

[heitorlessa]

Thanks again @huonw for taking the time, truly appreciate that. Hopefully, we will be able to solve this and other similar situations with typing_extensions in the future

[huonw]

Woohoo, thanks for the tweak and then merging 👍