diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c48b2f9..7f201f6 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3875,7 +3875,7 @@
     if( ret != 0 )
         goto exit;
 
-#if defined(MBEDTLS_DES_C)
+#if defined(PSA_WANT_KEY_TYPE_DES)
     if( slot->attr.type == PSA_KEY_TYPE_DES && key_bits == 128 )
     {
         /* Two-key Triple-DES is 3-key Triple-DES with K1=K3 */
@@ -5063,7 +5063,7 @@
     return( status );
 }
 
-#if defined(MBEDTLS_DES_C)
+#if defined(PSA_WANT_KEY_TYPE_DES)
 static void psa_des_set_key_parity( uint8_t *data, size_t data_size )
 {
     if( data_size >= 8 )
@@ -5073,7 +5073,7 @@
     if( data_size >= 24 )
         mbedtls_des_key_set_parity( data + 16 );
 }
-#endif /* MBEDTLS_DES_C */
+#endif /* PSA_WANT_KEY_TYPE_DES */
 
 static psa_status_t psa_generate_derived_key_internal(
     psa_key_slot_t *slot,
@@ -5095,10 +5095,10 @@
     status = psa_key_derivation_output_bytes( operation, data, bytes );
     if( status != PSA_SUCCESS )
         goto exit;
-#if defined(MBEDTLS_DES_C)
+#if defined(PSA_WANT_KEY_TYPE_DES)
     if( slot->attr.type == PSA_KEY_TYPE_DES )
         psa_des_set_key_parity( data, bytes );
-#endif /* MBEDTLS_DES_C */
+#endif /* PSA_WANT_KEY_TYPE_DES */
 
     status = psa_allocate_buffer_to_slot( slot, bytes );
     if( status != PSA_SUCCESS )
@@ -6029,7 +6029,7 @@
         if( status != PSA_SUCCESS )
             return( status );
 
-#if defined(MBEDTLS_DES_C)
+#if defined(PSA_WANT_KEY_TYPE_DES)
         if( type == PSA_KEY_TYPE_DES )
             psa_des_set_key_parity( key_buffer, key_buffer_size );
 #endif /* MBEDTLS_DES_C */