diff --git a/scripts/mbedtls_dev/test_generation.py b/scripts/mbedtls_dev/test_generation.py
index 55076dc..64b8fe4 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.
@@ -149,7 +148,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.