fix: streamlined regexes

Author: rubenfonseca

aws_lambda_powertools/event_handler/api_gateway.py

@@ -520,7 +520,7 @@ def __init__(
         cors: Optional[CORSConfig] = None,
         debug: Optional[bool] = None,
         serializer: Optional[Callable[[Dict], str]] = None,
-        strip_prefixes: Optional[List[str]] = None,
+        strip_prefixes: Optional[List[Union[str, Pattern]]] = None,
     ):
         """
         Parameters
@@ -534,9 +534,9 @@ def __init__(
             environment variable
         serializer : Callable, optional
             function to serialize `obj` to a JSON formatted `str`, by default json.dumps
-        strip_prefixes: List[str], optional
+        strip_prefixes: List[Union[str, Pattern]], optional
             optional list of prefixes to be removed from the request path before doing the routing. This is often used
-            with api gateways with multiple custom mappings.
+            with api gateways with multiple custom mappings. Each prefix can be a static string or a compiled regex
         """
         self._proxy_type = proxy_type
         self._dynamic_routes: List[Route] = []
@@ -713,10 +713,18 @@ def _remove_prefix(self, path: str) -> str:
             return path
 
         for prefix in self._strip_prefixes:
-            if path == prefix:
-                return "/"
-            path = re.sub(rf"^/?({prefix})/", r"/", path)
+            if isinstance(prefix, str):
+                if path == prefix:
+                    return "/"
+
+                if self._path_starts_with(path, prefix):
+                    return path[len(prefix) :]
+
+            if isinstance(prefix, Pattern):
+                path = re.sub(prefix, "", path)
 
+        # When using regexes, we might get into a point where everything is removed
+        # from the string, so we check if it's empty and change it accordingly.
         if not path:
             path = "/"
 

tests/functional/event_handler/test_api_gateway.py

@@ -1,5 +1,6 @@
 import base64
 import json
+import re
 import zlib
 from copy import deepcopy
 from decimal import Decimal
@@ -1077,6 +1078,26 @@ def foo():
     assert response["statusCode"] == 200
 
 
+@pytest.mark.parametrize(
+    "path",
+    [
+        pytest.param("/stg/foo", id="path matched pay prefix"),
+        pytest.param("/dev/foo", id="path matched pay prefix with multiple numbers"),
+        pytest.param("/foo", id="path does not start with any of the prefixes"),
+    ],
+)
+def test_remove_prefix_by_regex(path: str):
+    app = ApiGatewayResolver(strip_prefixes=[re.compile(r"/(dev|stg)")])
+
+    @app.get("/foo")
+    def foo():
+        ...
+
+    response = app({"httpMethod": "GET", "path": path}, None)
+
+    assert response["statusCode"] == 200
+
+
 @pytest.mark.parametrize(
     "prefix",
     [