Pass GenerationContext to customizeBeanRegistrationCodeFragments

Update `customizeBeanRegistrationCodeFragments` so that it can access the `GenerationContext`. See gh-28557
spring-projects · Jun 6, 2022 · 305055d · 305055d
1 parent 9a9c3ea
commit 305055d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/...ns/src/main/java/org/springframework/beans/factory/aot/BeanDefinitionMethodGenerator.java b/...ns/src/main/java/org/springframework/beans/factory/aot/BeanDefinitionMethodGenerator.java
@@ -88,7 +88,8 @@ class BeanDefinitionMethodGenerator {
 	MethodReference generateBeanDefinitionMethod(GenerationContext generationContext,
 			String featureNamePrefix, BeanRegistrationsCode beanRegistrationsCode) {
 
-		BeanRegistrationCodeFragments codeFragments = getCodeFragments(beanRegistrationsCode, featureNamePrefix);
+		BeanRegistrationCodeFragments codeFragments = getCodeFragments(generationContext,
+				beanRegistrationsCode, featureNamePrefix);
 		Class<?> target = codeFragments.getTarget(this.registeredBean,
 				this.constructorOrFactoryMethod);
 		if (!target.getName().startsWith("java.")) {
@@ -113,14 +114,14 @@ MethodReference generateBeanDefinitionMethod(GenerationContext generationContext
 				generatedMethod.getName().toString());
 	}
 
-	private BeanRegistrationCodeFragments getCodeFragments(
+	private BeanRegistrationCodeFragments getCodeFragments(GenerationContext generationContext,
 			BeanRegistrationsCode beanRegistrationsCode, String featureNamePrefix) {
 
 		BeanRegistrationCodeFragments codeFragments = new DefaultBeanRegistrationCodeFragments(
 				beanRegistrationsCode, this.registeredBean, this.methodGeneratorFactory,
 				featureNamePrefix);
 		for (BeanRegistrationAotContribution aotContribution : this.aotContributions) {
-			codeFragments = aotContribution.customizeBeanRegistrationCodeFragments(codeFragments);
+			codeFragments = aotContribution.customizeBeanRegistrationCodeFragments(generationContext, codeFragments);
 		}
 		return codeFragments;
 	}

diff --git a/.../src/main/java/org/springframework/beans/factory/aot/BeanRegistrationAotContribution.java b/.../src/main/java/org/springframework/beans/factory/aot/BeanRegistrationAotContribution.java
@@ -36,12 +36,13 @@ public interface BeanRegistrationAotContribution {
 	 * Customize the {@link BeanRegistrationCodeFragments} that will be used to
 	 * generate the bean registration code. Custom code fragments can be used if
 	 * default code generation isn't suitable.
+	 * @param generationContext the generation context
 	 * @param codeFragments the existing code fragments
 	 * @return the code fragments to use, may be the original instance or a
 	 * wrapper
 	 */
 	default BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
-			BeanRegistrationCodeFragments codeFragments) {
+			GenerationContext generationContext, BeanRegistrationCodeFragments codeFragments) {
 		return codeFragments;
 	}
 
@@ -70,7 +71,7 @@ static BeanRegistrationAotContribution ofBeanRegistrationCodeFragmentsCustomizer
 
 			@Override
 			public BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(
-					BeanRegistrationCodeFragments codeFragments) {
+					GenerationContext generationContext, BeanRegistrationCodeFragments codeFragments) {
 				return beanRegistrationCodeFragmentsCustomizer.apply(codeFragments);
 			}