Add type hinting for discord provider (#9773)

GitOrigin-RevId: a518801f8d5abe4ceb8b8678c27e6858f51f288a
kosteev · Sep 12, 2024 · 0699aec · 0699aec
1 parent 90ae740
commit 0699aec
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 22 deletions.
diff --git a/airflow/providers/discord/hooks/discord_webhook.py b/airflow/providers/discord/hooks/discord_webhook.py
@@ -18,6 +18,7 @@
 #
 import json
 import re
+from typing import Any, Dict, Optional
 
 from airflow.exceptions import AirflowException
 from airflow.providers.http.hooks.http import HttpHook
@@ -54,15 +55,15 @@ class DiscordWebhookHook(HttpHook):
     """
 
     def __init__(self,
-                 http_conn_id=None,
-                 webhook_endpoint=None,
-                 message="",
-                 username=None,
-                 avatar_url=None,
-                 tts=False,
-                 proxy=None,
+                 http_conn_id: Optional[str] = None,
+                 webhook_endpoint: Optional[str] = None,
+                 message: str = "",
+                 username: Optional[str] = None,
+                 avatar_url: Optional[str] = None,
+                 tts: bool = False,
+                 proxy: Optional[str] = None,
                  *args,
-                 **kwargs):
+                 **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.http_conn_id = http_conn_id
         self.webhook_endpoint = self._get_webhook_endpoint(http_conn_id, webhook_endpoint)
@@ -72,7 +73,7 @@ def __init__(self,
         self.tts = tts
         self.proxy = proxy
 
-    def _get_webhook_endpoint(self, http_conn_id, webhook_endpoint):
+    def _get_webhook_endpoint(self, http_conn_id: Optional[str], webhook_endpoint: Optional[str]) -> str:
         """
         Given a Discord http_conn_id, return the default webhook endpoint or override if a
         webhook_endpoint is manually supplied.
@@ -98,14 +99,14 @@ def _get_webhook_endpoint(self, http_conn_id, webhook_endpoint):
 
         return endpoint
 
-    def _build_discord_payload(self):
+    def _build_discord_payload(self) -> str:
         """
         Construct the Discord JSON payload. All relevant parameters are combined here
         to a valid Discord JSON payload.
 
         :return: Discord payload (str) to send
         """
-        payload = {}
+        payload: Dict[str, Any] = {}
 
         if self.username:
             payload['username'] = self.username
@@ -122,7 +123,7 @@ def _build_discord_payload(self):
 
         return json.dumps(payload)
 
-    def execute(self):
+    def execute(self) -> None:
         """
         Execute the Discord webhook call
         """

diff --git a/airflow/providers/discord/operators/discord_webhook.py b/airflow/providers/discord/operators/discord_webhook.py
@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+from typing import Dict, Optional
+
 from airflow.exceptions import AirflowException
 from airflow.providers.discord.hooks.discord_webhook import DiscordWebhookHook
 from airflow.providers.http.operators.http import SimpleHttpOperator
@@ -56,15 +58,15 @@ class DiscordWebhookOperator(SimpleHttpOperator):
 
     @apply_defaults
     def __init__(self,
-                 http_conn_id=None,
-                 webhook_endpoint=None,
-                 message="",
-                 username=None,
-                 avatar_url=None,
-                 tts=False,
-                 proxy=None,
+                 http_conn_id: Optional[str] = None,
+                 webhook_endpoint: Optional[str] = None,
+                 message: str = "",
+                 username: Optional[str] = None,
+                 avatar_url: Optional[str] = None,
+                 tts: bool = False,
+                 proxy: Optional[str] = None,
                  *args,
-                 **kwargs):
+                 **kwargs) -> None:
         super().__init__(endpoint=webhook_endpoint,
                          *args,
                          **kwargs)
@@ -79,9 +81,9 @@ def __init__(self,
         self.avatar_url = avatar_url
         self.tts = tts
         self.proxy = proxy
-        self.hook = None
+        self.hook: Optional[DiscordWebhookHook] = None
 
-    def execute(self, context):
+    def execute(self, context: Dict) -> None:
         """
         Call the DiscordWebhookHook to post message
         """