diff --git a/ddtrace/appsec/_utils.py b/ddtrace/appsec/_utils.py index 5f0fee4f54f..5ea7ac1475f 100644 --- a/ddtrace/appsec/_utils.py +++ b/ddtrace/appsec/_utils.py @@ -173,14 +173,26 @@ def __init__( self.location = location.replace(APPSEC.SECURITY_RESPONSE_ID, security_response_id) self.content_type: str = "application/json" - def get(self, method_name: str, default: Any = None) -> Any: + def get(self, key: str, default: Any = None) -> Union[str, int]: """ Dictionary-like get method for backward compatibility with Lambda integration. Returns the attribute value if it exists, otherwise returns the default value. This allows Block_config to be used in contexts that expect dictionary-like access. """ - return getattr(self, method_name, default) + if key == "content-type": + key = "content_type" + return getattr(self, key, default) + + def __getitem__(self, key: str) -> Optional[Union[str, int]]: + if key == "content-type": + key = "content_type" + return getattr(self, key, None) + + def __contains__(self, key: str) -> bool: + if key == "content-type": + key = "content_type" + return bool(getattr(self, key, None)) class Telemetry_result: diff --git a/ddtrace/internal/utils/__init__.py b/ddtrace/internal/utils/__init__.py index e5bf0862851..fbaac5a684b 100644 --- a/ddtrace/internal/utils/__init__.py +++ b/ddtrace/internal/utils/__init__.py @@ -87,6 +87,12 @@ class Block_config(Protocol): location: str content_type: str + def get(self, key: str, default: Any = None) -> Union[str, int]: ... + + def __getitem__(self, key: str) -> Optional[Union[str, int]]: ... + + def __contains__(self, key: str) -> bool: ... + def get_blocked() -> Optional[Block_config]: # local import to avoid circular dependency