Skip to content

Commit 8fd5887

Browse files
chore(aap): improve internal api protocol (#15332)
## Description Add missing method in BlockConfig protocol use in our internal AAP API. Add missing method in BlockConfig class. ## Additional Notes This should only impact lambda integration. --------- Co-authored-by: Alberto Vara <[email protected]>
1 parent 38155b3 commit 8fd5887

File tree

2 files changed

+20
-2
lines changed

2 files changed

+20
-2
lines changed

ddtrace/appsec/_utils.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,14 +173,26 @@ def __init__(
173173
self.location = location.replace(APPSEC.SECURITY_RESPONSE_ID, security_response_id)
174174
self.content_type: str = "application/json"
175175

176-
def get(self, method_name: str, default: Any = None) -> Any:
176+
def get(self, key: str, default: Any = None) -> Union[str, int]:
177177
"""
178178
Dictionary-like get method for backward compatibility with Lambda integration.
179179
180180
Returns the attribute value if it exists, otherwise returns the default value.
181181
This allows Block_config to be used in contexts that expect dictionary-like access.
182182
"""
183-
return getattr(self, method_name, default)
183+
if key == "content-type":
184+
key = "content_type"
185+
return getattr(self, key, default)
186+
187+
def __getitem__(self, key: str) -> Optional[Union[str, int]]:
188+
if key == "content-type":
189+
key = "content_type"
190+
return getattr(self, key, None)
191+
192+
def __contains__(self, key: str) -> bool:
193+
if key == "content-type":
194+
key = "content_type"
195+
return bool(getattr(self, key, None))
184196

185197

186198
class Telemetry_result:

ddtrace/internal/utils/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,12 @@ class Block_config(Protocol):
8787
location: str
8888
content_type: str
8989

90+
def get(self, key: str, default: Any = None) -> Union[str, int]: ...
91+
92+
def __getitem__(self, key: str) -> Optional[Union[str, int]]: ...
93+
94+
def __contains__(self, key: str) -> bool: ...
95+
9096

9197
def get_blocked() -> Optional[Block_config]:
9298
# local import to avoid circular dependency

0 commit comments

Comments
 (0)