[airflow] branch master updated: Add typing annotations to Segment provider (#10120)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/master by this push:
     new 85cc2a6  Add typing annotations to Segment provider (#10120)
85cc2a6 is described below

commit 85cc2a605e82c6ad26f310e989b365d9d490f2f3
Author: chipmyersjr <ch...@gmail.com>
AuthorDate: Mon Aug 3 00:53:35 2020 -0700

    Add typing annotations to Segment provider (#10120)
---
 airflow/providers/segment/hooks/segment.py              | 10 +++++-----
 .../providers/segment/operators/segment_track_event.py  | 17 +++++++++--------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/airflow/providers/segment/hooks/segment.py b/airflow/providers/segment/hooks/segment.py
index 940621d..2ff223e 100644
--- a/airflow/providers/segment/hooks/segment.py
+++ b/airflow/providers/segment/hooks/segment.py
@@ -54,11 +54,11 @@ class SegmentHook(BaseHook):
     """
     def __init__(
             self,
-            segment_conn_id='segment_default',
-            segment_debug_mode=False,
+            segment_conn_id: str = 'segment_default',
+            segment_debug_mode: bool = False,
             *args,
             **kwargs
-    ):
+    ) -> None:
         super().__init__()
         self.segment_conn_id = segment_conn_id
         self.segment_debug_mode = segment_debug_mode
@@ -72,7 +72,7 @@ class SegmentHook(BaseHook):
         if self.write_key is None:
             raise AirflowException('No Segment write key provided')
 
-    def get_conn(self):
+    def get_conn(self) -> analytics:
         self.log.info('Setting write key for Segment analytics connection')
         analytics.debug = self.segment_debug_mode
         if self.segment_debug_mode:
@@ -81,7 +81,7 @@ class SegmentHook(BaseHook):
         analytics.write_key = self.write_key
         return analytics
 
-    def on_error(self, error, items):
+    def on_error(self, error: str, items: str) -> None:
         """
         Handles error callbacks when using Segment with segment_debug_mode set to True
         """
diff --git a/airflow/providers/segment/operators/segment_track_event.py b/airflow/providers/segment/operators/segment_track_event.py
index 12adf17..06025a4 100644
--- a/airflow/providers/segment/operators/segment_track_event.py
+++ b/airflow/providers/segment/operators/segment_track_event.py
@@ -15,6 +15,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from typing import Dict, Optional
 
 from airflow.models import BaseOperator
 from airflow.providers.segment.hooks.segment import SegmentHook
@@ -42,12 +43,12 @@ class SegmentTrackEventOperator(BaseOperator):
 
     @apply_defaults
     def __init__(self,
-                 user_id,
-                 event,
-                 properties=None,
-                 segment_conn_id='segment_default',
-                 segment_debug_mode=False,
-                 **kwargs):
+                 user_id: str,
+                 event: str,
+                 properties: Optional[dict] = None,
+                 segment_conn_id: str = 'segment_default',
+                 segment_debug_mode: bool = False,
+                 **kwargs) -> None:
         super().__init__(**kwargs)
         self.user_id = user_id
         self.event = event
@@ -56,7 +57,7 @@ class SegmentTrackEventOperator(BaseOperator):
         self.segment_debug_mode = segment_debug_mode
         self.segment_conn_id = segment_conn_id
 
-    def execute(self, context):
+    def execute(self, context: Dict) -> None:
         hook = SegmentHook(segment_conn_id=self.segment_conn_id,
                            segment_debug_mode=self.segment_debug_mode)
 
@@ -65,7 +66,7 @@ class SegmentTrackEventOperator(BaseOperator):
             self.event, self.user_id, self.properties)
 
         # pylint: disable=no-member
-        hook.track(
+        hook.track(  # type: ignore
             user_id=self.user_id,
             event=self.event,
             properties=self.properties)