c_wrapper_generator: Refactored _write_prologue()

Signed-off-by: Minos Galanakis <minos.galanakis@arm.com>

scripts/mbedtls_framework/c_wrapper_generator.py

@@ -46,6 +46,8 @@ class Base:
     # Suffix appended to the function's name to form the wrapper name.
     _WRAPPER_NAME_SUFFIX = '_wrap'
 
+    _INCLUDES = ['<mbedtls/build_info.h>']
+
     # Functions with one of these qualifiers are skipped.
     _SKIP_FUNCTION_WITH_QUALIFIERS = frozenset(['inline', 'static'])
 
@@ -65,26 +67,29 @@ class Base:
 
         This includes a description comment and some include directives.
         """
-        out.write("""/* Automatically generated by {}, do not edit! */
+        prologue = ['/* Automatically generated by {}, do not edit! */'.format(self.program_name),
+                    '']
+        prologue += ['/* Copyright The Mbed TLS Contributors',
+                     ' * SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later',
+                     ' */',
+                     '']
 
-/* Copyright The Mbed TLS Contributors
- * SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
- */
-"""
-                  .format(self.program_name))
         if header:
-            out.write("""
-#ifndef {guard}
-#define {guard}
+            prologue += ['#ifndef {}'.format(self.header_guard),
+                         '#define {}'.format(self.header_guard),
+                         '',
+                         '#ifdef __cplusplus',
+                         'extern "C" {',
+                         '#endif']
 
-#ifdef __cplusplus
-extern "C" {{
-#endif
-"""
-                      .format(guard=self.header_guard))
-        out.write("""
-#include <mbedtls/build_info.h>
-""")
+        for include in self._INCLUDES:
+            prologue.append('#include {}'.format(include))
+
+        # Make certain there is an empty line
+        if prologue[-1] != '':
+            prologue += (['', ''])
+
+        out.write('\n'.join(prologue))
 
     def _write_epilogue(self, out: typing_util.Writable, header: bool) -> None:
         """Write the epilogue of a C file.

scripts/mbedtls_framework/code_wrapper/psa_test_wrapper.py

@@ -31,16 +31,10 @@ class PSATestWrapper(PSAWrapper):
     _WRAPPER_NAME_PREFIX = 'mbedtls_test_wrap_'
     _WRAPPER_NAME_SUFFIX = ''
 
-    __PROLOGUE__ = """
-        #if {}
-
-        #include <psa/crypto.h>
-
-        #include <test/memory.h>
-        #include <test/psa_crypto_helpers.h>
-        #include <test/psa_test_wrappers.h>
-        """
-
+    _PSA_WRAPPER_INCLUDES = ['<psa/crypto.h>',
+                             '<test/memory.h>',
+                             '<test/psa_crypto_helpers.h>',
+                             '<test/psa_test_wrappers.h>']
 
 class PSALoggingTestWrapper(PSATestWrapper, PSALoggingWrapper):
     """Generate a C source file containing wrapper functions that log PSA Crypto API calls."""

scripts/mbedtls_framework/code_wrapper/psa_wrapper.py

@@ -64,15 +64,7 @@ class PSAWrapper(c_wrapper_generator.Base):
     _WRAPPER_NAME_PREFIX = 'mbedtls_test_wrap_'
     _WRAPPER_NAME_SUFFIX = ''
 
-    __PROLOGUE__ = """
-        #if {}
-
-        #include <psa/crypto.h>
-        """
-
-    __EPILOGUE__ = """
-        #endif /* {} */
-        """
+    _PSA_WRAPPER_INCLUDES = ['<psa/crypto.h>']
 
     def __init__(self,
                  output_h_f: str,
@@ -111,7 +103,7 @@ class PSAWrapper(c_wrapper_generator.Base):
         """ Return the estimated path in relationship to the.
             mbedtls_root. The method allows overriding the targetted sub-directory.
             Currently the default is set to mbedtls_root/include/psa """
-        
+
         # Temporary, while Mbed TLS does not just rely on the TF-PSA-Crypto
         # build system to build its crypto library. When it does, the first
         # case can just be removed.
@@ -197,8 +189,19 @@ class PSAWrapper(c_wrapper_generator.Base):
 
     def _write_prologue(self, out: typing_util.Writable, header: bool) -> None:
         super()._write_prologue(out, header)
+
+        prologue = []
         if self._CPP_GUARDS:
-            out.write(dedent(self.__PROLOGUE__).format(self._CPP_GUARDS))
+            prologue.append("#if {}".format(self._CPP_GUARDS))
+            prologue.append('')
+
+        for include in self._PSA_WRAPPER_INCLUDES:
+            prologue.append("#include {}".format(include))
+
+        if prologue[-1] != '':
+            prologue.append('')
+
+        out.write("\n".join(prologue))
 
     def _write_epilogue(self, out: typing_util.Writable, header: bool) -> None:
         if self._CPP_GUARDS: