From e3ad22ecf256a6d32bb4e417da01b611ffb9ed94 Mon Sep 17 00:00:00 2001
From: Werner Lewis <werner.lewis@arm.com>
Date: Thu, 25 Aug 2022 10:02:06 +0100
Subject: [PATCH] Fix TARGET types and code style

Signed-off-by: Werner Lewis <werner.lewis@arm.com>
---
 scripts/mbedtls_dev/test_generation.py |  3 +--
 tests/scripts/generate_bignum_tests.py | 12 ++++++------
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/scripts/mbedtls_dev/test_generation.py b/scripts/mbedtls_dev/test_generation.py
index a90547349..11c085f6b 100644
--- a/scripts/mbedtls_dev/test_generation.py
+++ b/scripts/mbedtls_dev/test_generation.py
@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 """Common test generation classes and main function.
 
 These are used both by generate_psa_tests.py and generate_bignum_tests.py.
@@ -150,7 +149,7 @@ class TestGenerator:
 
     # Note that targets whose names contain 'test_format' have their content
     # validated by `abi_check.py`.
-    TARGETS = {} # type: Dict[str, Callable[..., test_case.TestCase]]
+    TARGETS = {} # type: Dict[str, Callable[..., Iterable[test_case.TestCase]]]
 
     def generate_target(self, name: str, *target_args) -> None:
         """Generate cases and write to data file for a target.
diff --git a/tests/scripts/generate_bignum_tests.py b/tests/scripts/generate_bignum_tests.py
index aa7e131a1..c57f197a2 100755
--- a/tests/scripts/generate_bignum_tests.py
+++ b/tests/scripts/generate_bignum_tests.py
@@ -49,7 +49,7 @@ import itertools
 import sys
 
 from abc import ABCMeta, abstractmethod
-from typing import Callable, Dict, Iterator, List, Tuple, TypeVar, cast
+from typing import Callable, Dict, Iterable, Iterator, List, Tuple, TypeVar, cast
 
 import scripts_path # pylint: disable=unused-import
 from mbedtls_dev import test_case
@@ -206,10 +206,10 @@ class BignumAdd(BignumOperation):
     input_cases = cast(
         List[Tuple[str, str]],
         list(itertools.combinations(
-        [
-            "1c67967269c6", "9cde3",
-            "-1c67967269c6", "-9cde3",
-        ], 2
+            [
+                "1c67967269c6", "9cde3",
+                "-1c67967269c6", "-9cde3",
+            ], 2
         ))
     )
 
@@ -226,7 +226,7 @@ class BignumTestGenerator(test_generation.TestGenerator):
     TARGETS = {
         subclass.target_basename: subclass.generate_tests for subclass in
         test_generation.BaseTarget.__subclasses__()
-    } # type: Dict[str, Callable[[], test_case.TestCase]]
+    } # type: Dict[str, Callable[[], Iterable[test_case.TestCase]]]
 
 if __name__ == '__main__':
     test_generation.main(sys.argv[1:], BignumTestGenerator)