From 10ab267afb763a8776495de69eda0e4d6a4d3a79 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Wed, 10 Mar 2021 00:59:53 +0100
Subject: [PATCH] Add some type annotations

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
---
 scripts/mbedtls_dev/macro_collector.py   | 27 ++++++++--------
 tests/scripts/test_psa_constant_names.py | 39 ++++++++++++------------
 2 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/scripts/mbedtls_dev/macro_collector.py b/scripts/mbedtls_dev/macro_collector.py
index b98e40e5fc..7ebd8f7c3f 100644
--- a/scripts/mbedtls_dev/macro_collector.py
+++ b/scripts/mbedtls_dev/macro_collector.py
@@ -17,12 +17,13 @@
 # limitations under the License.
 
 import re
+from typing import Dict, Set
 
 class PSAMacroCollector:
     """Collect PSA crypto macro definitions from C header files.
     """
 
-    def __init__(self, include_intermediate=False):
+    def __init__(self, include_intermediate: bool = False) -> None:
         """Set up an object to collect PSA macro definitions.
 
         Call the read_file method of the constructed object on each header file.
@@ -31,19 +32,19 @@ class PSAMacroCollector:
           PSA_XXX_BASE that do not designate semantic values.
         """
         self.include_intermediate = include_intermediate
-        self.statuses = set()
-        self.key_types = set()
-        self.key_types_from_curve = {}
-        self.key_types_from_group = {}
-        self.ecc_curves = set()
-        self.dh_groups = set()
-        self.algorithms = set()
-        self.hash_algorithms = set()
-        self.ka_algorithms = set()
-        self.algorithms_from_hash = {}
-        self.key_usages = set()
+        self.statuses = set() #type: Set[str]
+        self.key_types = set() #type: Set[str]
+        self.key_types_from_curve = {} #type: Dict[str, str]
+        self.key_types_from_group = {} #type: Dict[str, str]
+        self.ecc_curves = set() #type: Set[str]
+        self.dh_groups = set() #type: Set[str]
+        self.algorithms = set() #type: Set[str]
+        self.hash_algorithms = set() #type: Set[str]
+        self.ka_algorithms = set() #type: Set[str]
+        self.algorithms_from_hash = {} #type: Dict[str, str]
+        self.key_usages = set() #type: Set[str]
 
-    def is_internal_name(self, name):
+    def is_internal_name(self, name: str) -> bool:
         """Whether this is an internal macro. Internal macros will be skipped."""
         if not self.include_intermediate:
             if name.endswith('_BASE') or name.endswith('_NONE'):
diff --git a/tests/scripts/test_psa_constant_names.py b/tests/scripts/test_psa_constant_names.py
index 9795e05c07..21a2e8d254 100755
--- a/tests/scripts/test_psa_constant_names.py
+++ b/tests/scripts/test_psa_constant_names.py
@@ -29,6 +29,7 @@ import os
 import re
 import subprocess
 import sys
+from typing import Dict, Iterable, Iterator, List, Set
 
 import scripts_path # pylint: disable=unused-import
 from mbedtls_dev import c_build_helper
@@ -46,31 +47,31 @@ class PSAMacroEnumerator:
     enumerated here.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Set up an empty set of known constructor macros.
         """
-        self.statuses = set()
-        self.algorithms = set()
-        self.ecc_curves = set()
-        self.dh_groups = set()
-        self.key_types = set()
-        self.key_usage_flags = set()
-        self.hash_algorithms = set()
-        self.mac_algorithms = set()
-        self.ka_algorithms = set()
-        self.kdf_algorithms = set()
-        self.aead_algorithms = set()
+        self.statuses = set() #type: Set[str]
+        self.algorithms = set() #type: Set[str]
+        self.ecc_curves = set() #type: Set[str]
+        self.dh_groups = set() #type: Set[str]
+        self.key_types = set() #type: Set[str]
+        self.key_usage_flags = set() #type: Set[str]
+        self.hash_algorithms = set() #type: Set[str]
+        self.mac_algorithms = set() #type: Set[str]
+        self.ka_algorithms = set() #type: Set[str]
+        self.kdf_algorithms = set() #type: Set[str]
+        self.aead_algorithms = set() #type: Set[str]
         # macro name -> list of argument names
-        self.argspecs = {}
+        self.argspecs = {} #type: Dict[str, List[str]]
         # argument name -> list of values
         self.arguments_for = {
             'mac_length': [],
             'min_mac_length': [],
             'tag_length': [],
             'min_tag_length': [],
-        }
+        } #type: Dict[str, List[str]]
 
-    def gather_arguments(self):
+    def gather_arguments(self) -> None:
         """Populate the list of values for macro arguments.
 
         Call this after parsing all the inputs.
@@ -84,16 +85,16 @@ class PSAMacroEnumerator:
         self.arguments_for['group'] = sorted(self.dh_groups)
 
     @staticmethod
-    def _format_arguments(name, arguments):
+    def _format_arguments(name: str, arguments: Iterable[str]) -> str:
         """Format a macro call with arguments.."""
         return name + '(' + ', '.join(arguments) + ')'
 
     _argument_split_re = re.compile(r' *, *')
     @classmethod
-    def _argument_split(cls, arguments):
+    def _argument_split(cls, arguments: str) -> List[str]:
         return re.split(cls._argument_split_re, arguments)
 
-    def distribute_arguments(self, name):
+    def distribute_arguments(self, name: str) -> Iterator[str]:
         """Generate macro calls with each tested argument set.
 
         If name is a macro without arguments, just yield "name".
@@ -123,7 +124,7 @@ class PSAMacroEnumerator:
         except BaseException as e:
             raise Exception('distribute_arguments({})'.format(name)) from e
 
-    def generate_expressions(self, names):
+    def generate_expressions(self, names: Iterable[str]) -> Iterator[str]:
         """Generate expressions covering values constructed from the given names.
 
         `names` can be any iterable collection of macro names.