diff --git a/scripts/code_style.py b/scripts/code_style.py
index 63cc6dc7a..ef2750819 100755
--- a/scripts/code_style.py
+++ b/scripts/code_style.py
@@ -41,16 +41,29 @@ def list_generated_files() -> FrozenSet[str]:
     would conform to the code style, but this would be difficult, especially
     with respect to the placement of line breaks in long logical lines.
     """
-    # Parse check-generated-files.sh to get an up-to-date list of
-    # generated files. Read the file rather than calling it so that
-    # this script only depends on Git, Python and uncrustify, and not other
-    # tools such as sh or grep which might not be available on Windows.
-    # This introduces a limitation: check-generated-files.sh must have
-    # the expected format and must list the files explicitly, not through
-    # wildcards or command substitution.
-    content = open(CHECK_GENERATED_FILES, encoding="utf-8").read()
-    checks = re.findall(CHECK_CALL_RE, content)
-    return frozenset(word for s in checks for word in s.split())
+    if build_tree.is_mbedtls_3_6():
+        # Parse check-generated-files.sh to get an up-to-date list of
+        # generated files. Read the file rather than calling it so that
+        # this script only depends on Git, Python and uncrustify, and not other
+        # tools such as sh or grep which might not be available on Windows.
+        # This introduces a limitation: check-generated-files.sh must have
+        # the expected format and must list the files explicitly, not through
+        # wildcards or command substitution.
+        content = open(CHECK_GENERATED_FILES, encoding="utf-8").read()
+        checks = re.findall(CHECK_CALL_RE, content)
+        return frozenset(word for s in checks for word in s.split())
+    else:
+        output = subprocess.check_output(["framework/scripts/make_generated_files.py",
+                                          "--list"], universal_newlines=True)
+        # psa_test_wrappers.[hc], generated by generate_psa_wrappers.py, are
+        # currently committed and unknown to make_generated_files.py. Add them
+        # here to the list of generated file as we do not want to check their
+        # coding style.
+        if build_tree.looks_like_tf_psa_crypto_root("."):
+            output += "tests/include/test/psa_test_wrappers.h\n"
+            output += "tests/src/psa_test_wrappers.c"
+
+        return frozenset(line for line in output.splitlines())
 
 # Check for comment string indicating an auto-generated file
 AUTOGEN_RE = re.compile(r"Warning[ :-]+This file is (now )?auto[ -]?generated",
@@ -72,7 +85,6 @@ def get_src_files(since: Optional[str]) -> List[str]:
     """
     file_patterns = ["*.[hc]",
                      "tests/suites/*.function",
-                     "tf-psa-crypto/tests/suites/*.function",
                      "scripts/data_files/*.fmt"]
     output = subprocess.check_output(["git", "ls-files"] + file_patterns,
                                      universal_newlines=True)
@@ -140,8 +152,8 @@ def get_src_files(since: Optional[str]) -> List[str]:
                              is_file_autogenerated(filename))]
     else:
         src_files = [filename for filename in src_files
-                     if not (filename.startswith("tf-psa-crypto/drivers/everest/") or
-                             filename.startswith("tf-psa-crypto/drivers/p256-m/") or
+                     if not (filename.startswith("drivers/everest/") or
+                             filename.startswith("drivers/p256-m/") or
                              filename in generated_files or
                              is_file_autogenerated(filename))]
     return src_files