From 13362f536e9d8a092ec20dcb5abe7a0b86407f45 Mon Sep 17 00:00:00 2001
From: Johnson Shi <Johnson.Shi@microsoft.com>
Date: Fri, 28 Aug 2020 08:23:59 -0700
Subject: Add method type hints for Azure helper (#540)

This reverts commit 8d25d5e6fac39ab3319ec5d37d23196429fb0c95.
---
 cloudinit/sources/helpers/azure.py | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

(limited to 'cloudinit/sources/helpers/azure.py')

diff --git a/cloudinit/sources/helpers/azure.py b/cloudinit/sources/helpers/azure.py
index b968a96f..507f6ac8 100755
--- a/cloudinit/sources/helpers/azure.py
+++ b/cloudinit/sources/helpers/azure.py
@@ -288,7 +288,8 @@ class InvalidGoalStateXMLException(Exception):
 
 class GoalState:
 
-    def __init__(self, unparsed_xml, azure_endpoint_client):
+    def __init__(self, unparsed_xml: str,
+                 azure_endpoint_client: AzureEndpointHttpClient) -> None:
         """Parses a GoalState XML string and returns a GoalState object.
 
         @param unparsed_xml: string representing a GoalState XML.
@@ -478,7 +479,10 @@ class GoalStateHealthReporter:
 
     PROVISIONING_SUCCESS_STATUS = 'Ready'
 
-    def __init__(self, goal_state, azure_endpoint_client, endpoint):
+    def __init__(
+            self, goal_state: GoalState,
+            azure_endpoint_client: AzureEndpointHttpClient,
+            endpoint: str) -> None:
         """Creates instance that will report provisioning status to an endpoint
 
         @param goal_state: An instance of class GoalState that contains
@@ -495,7 +499,7 @@ class GoalStateHealthReporter:
         self._endpoint = endpoint
 
     @azure_ds_telemetry_reporter
-    def send_ready_signal(self):
+    def send_ready_signal(self) -> None:
         document = self.build_report(
             incarnation=self._goal_state.incarnation,
             container_id=self._goal_state.container_id,
@@ -513,8 +517,8 @@ class GoalStateHealthReporter:
         LOG.info('Reported ready to Azure fabric.')
 
     def build_report(
-            self, incarnation, container_id, instance_id,
-            status, substatus=None, description=None):
+            self, incarnation: str, container_id: str, instance_id: str,
+            status: str, substatus=None, description=None) -> str:
         health_detail = ''
         if substatus is not None:
             health_detail = self.HEALTH_DETAIL_SUBSECTION_XML_TEMPLATE.format(
@@ -530,7 +534,7 @@ class GoalStateHealthReporter:
         return health_report
 
     @azure_ds_telemetry_reporter
-    def _post_health_report(self, document):
+    def _post_health_report(self, document: str) -> None:
         push_log_to_kvp()
 
         # Whenever report_diagnostic_event(diagnostic_msg) is invoked in code,
@@ -726,7 +730,7 @@ class WALinuxAgentShim:
         return endpoint_ip_address
 
     @azure_ds_telemetry_reporter
-    def register_with_azure_and_fetch_data(self, pubkey_info=None):
+    def register_with_azure_and_fetch_data(self, pubkey_info=None) -> dict:
         """Gets the VM's GoalState from Azure, uses the GoalState information
         to report ready/send the ready signal/provisioning complete signal to
         Azure, and then uses pubkey_info to filter and obtain the user's
@@ -750,7 +754,7 @@ class WALinuxAgentShim:
         return {'public-keys': ssh_keys}
 
     @azure_ds_telemetry_reporter
-    def _fetch_goal_state_from_azure(self):
+    def _fetch_goal_state_from_azure(self) -> GoalState:
         """Fetches the GoalState XML from the Azure endpoint, parses the XML,
         and returns a GoalState object.
 
@@ -760,7 +764,7 @@ class WALinuxAgentShim:
         return self._parse_raw_goal_state_xml(unparsed_goal_state_xml)
 
     @azure_ds_telemetry_reporter
-    def _get_raw_goal_state_xml_from_azure(self):
+    def _get_raw_goal_state_xml_from_azure(self) -> str:
         """Fetches the GoalState XML from the Azure endpoint and returns
         the XML as a string.
 
@@ -780,7 +784,8 @@ class WALinuxAgentShim:
         return response.contents
 
     @azure_ds_telemetry_reporter
-    def _parse_raw_goal_state_xml(self, unparsed_goal_state_xml):
+    def _parse_raw_goal_state_xml(
+            self, unparsed_goal_state_xml: str) -> GoalState:
         """Parses a GoalState XML string and returns a GoalState object.
 
         @param unparsed_goal_state_xml: GoalState XML string
@@ -803,7 +808,8 @@ class WALinuxAgentShim:
         return goal_state
 
     @azure_ds_telemetry_reporter
-    def _get_user_pubkeys(self, goal_state, pubkey_info):
+    def _get_user_pubkeys(
+            self, goal_state: GoalState, pubkey_info: list) -> list:
         """Gets and filters the VM admin user's authorized pubkeys.
 
         The admin user in this case is the username specified as "admin"
@@ -838,7 +844,7 @@ class WALinuxAgentShim:
         return ssh_keys
 
     @staticmethod
-    def _filter_pubkeys(keys_by_fingerprint, pubkey_info):
+    def _filter_pubkeys(keys_by_fingerprint: dict, pubkey_info: list) -> list:
         """ Filter and return only the user's actual pubkeys.
 
         @param keys_by_fingerprint: pubkey fingerprint -> pubkey value dict
-- 
cgit v1.2.3