diff --git a/scripts/mbedtls_dev/test_generation.py b/scripts/mbedtls_dev/test_generation.py
index a905473..11c085f 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 @@
 
     # 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 aa7e131..c57f197 100755
--- a/tests/scripts/generate_bignum_tests.py
+++ b/tests/scripts/generate_bignum_tests.py
@@ -49,7 +49,7 @@
 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 @@
     input_cases = cast(
         List[Tuple[str, str]],
         list(itertools.combinations(
-        [
-            "1c67967269c6", "9cde3",
-            "-1c67967269c6", "-9cde3",
-        ], 2
+            [
+                "1c67967269c6", "9cde3",
+                "-1c67967269c6", "-9cde3",
+            ], 2
         ))
     )
 
@@ -226,7 +226,7 @@
     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)