From ad45da845da28d7e421f6b473efb787376a37162 Mon Sep 17 00:00:00 2001
From: Minos Galanakis <minos.galanakis@arm.com>
Date: Thu, 13 Jun 2024 11:41:29 +0100
Subject: [PATCH] code_wrappers: Updated import statements

Signed-off-by: Minos Galanakis <minos.galanakis@arm.com>
---
 scripts/generate_psa_wrappers.py                       | 10 +++-------
 .../mbedtls_framework/code_wrapper/psa_test_wrapper.py |  4 ++--
 scripts/mbedtls_framework/code_wrapper/psa_wrapper.py  |  2 +-
 3 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/scripts/generate_psa_wrappers.py b/scripts/generate_psa_wrappers.py
index 8ed10ed40..35602da9b 100755
--- a/scripts/generate_psa_wrappers.py
+++ b/scripts/generate_psa_wrappers.py
@@ -5,12 +5,8 @@
 # Copyright The Mbed TLS Contributors
 # SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
 
-### WARNING: the code in this file has not been extensively reviewed yet.
-### We do not think it is harmful, but it may be below our normal standards
-### for robustness and maintainability.
-
 import argparse
-from mbedtls_framework.code_wrapper import *
+from mbedtls_framework.code_wrapper.psa_test_wrapper import PSATestWrapper, PSALoggingTestWrapper
 
 DEFAULT_C_OUTPUT_FILE_NAME = 'tests/src/psa_test_wrappers.c'
 DEFAULT_H_OUTPUT_FILE_NAME = 'tests/include/test/psa_test_wrappers.h'
@@ -32,11 +28,11 @@ if __name__ == '__main__':
                               .format(DEFAULT_H_OUTPUT_FILE_NAME)))
     options = parser.parse_args() #pylint: disable=invalid-name
     if options.log:
-        generator = psa_test_wrapper.PSALoggingTestWrapper(DEFAULT_H_OUTPUT_FILE_NAME, #pylint: disable=invalid-name
+        generator = PSALoggingTestWrapper(DEFAULT_H_OUTPUT_FILE_NAME, #pylint: disable=invalid-name
                                           DEFAULT_C_OUTPUT_FILE_NAME,
                                           options.log) #type: PSATestWrapper
     else:
-        generator = psa_test_wrapper.PSATestWrapper(DEFAULT_H_OUTPUT_FILE_NAME, #pylint: disable=invalid-name
+        generator = PSATestWrapper(DEFAULT_H_OUTPUT_FILE_NAME, #pylint: disable=invalid-name
                                    DEFAULT_C_OUTPUT_FILE_NAME)
 
     if options.output_h:
diff --git a/scripts/mbedtls_framework/code_wrapper/psa_test_wrapper.py b/scripts/mbedtls_framework/code_wrapper/psa_test_wrapper.py
index 8feec58ca..222f8f141 100755
--- a/scripts/mbedtls_framework/code_wrapper/psa_test_wrapper.py
+++ b/scripts/mbedtls_framework/code_wrapper/psa_test_wrapper.py
@@ -19,8 +19,8 @@ from mbedtls_framework import c_parsing_helper
 from mbedtls_framework import c_wrapper_generator
 from mbedtls_framework import typing_util
 
-from psa_buffer import BufferParameter
-from psa_wrapper import PSAWrapper, PSALoggingWrapper, DEFAULTS
+from mbedtls_framework.code_wrapper.psa_buffer import BufferParameter
+from  mbedtls_framework.code_wrapper.psa_wrapper import PSAWrapper, PSALoggingWrapper, DEFAULTS
 
 class PSATestWrapper(PSAWrapper):
     """Generate a C source file containing wrapper functions for PSA Crypto API calls."""
diff --git a/scripts/mbedtls_framework/code_wrapper/psa_wrapper.py b/scripts/mbedtls_framework/code_wrapper/psa_wrapper.py
index 19ddb97a8..be5415748 100755
--- a/scripts/mbedtls_framework/code_wrapper/psa_wrapper.py
+++ b/scripts/mbedtls_framework/code_wrapper/psa_wrapper.py
@@ -15,7 +15,7 @@ from mbedtls_framework import c_parsing_helper
 from mbedtls_framework import c_wrapper_generator
 from mbedtls_framework import typing_util
 
-from psa_buffer import BufferParameter
+from mbedtls_framework.code_wrapper.psa_buffer import BufferParameter
 from textwrap import dedent
 
 DEFAULTS = {