Polishing.

Refine AotContext construction introducing a builder. Refine TypeCollector construction.

See #3474
Original pull request: #3475

Author: mp911de

src/main/java/org/springframework/data/aot/AotContext.java

@@ -36,6 +36,7 @@
 import org.springframework.core.env.Environment;
 import org.springframework.core.env.EnvironmentCapable;
 import org.springframework.core.env.StandardEnvironment;
+import org.springframework.data.util.TypeCollector;
 import org.springframework.data.util.TypeScanner;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
@@ -67,7 +68,6 @@ public interface AotContext extends EnvironmentCapable {
 	static AotContext from(BeanFactory beanFactory) {
 
 		Assert.notNull(beanFactory, "BeanFactory must not be null");
-
 		return new DefaultAotContext(beanFactory, new StandardEnvironment());
 	}
 
@@ -87,6 +87,13 @@ static AotContext from(BeanFactory beanFactory, Environment environment) {
 		return new DefaultAotContext(beanFactory, environment);
 	}
 
+	/**
+	 * Create a {@code AotContext} builder.
+	 */
+	public static AotContextBuilder builder() {
+		return new DefaultAotContextBuilder();
+	}
+
 	/**
 	 * Checks if repository code generation is enabled for a given module by checking environment variables for general
 	 * enablement ({@link #GENERATED_REPOSITORIES_ENABLED}) and store-specific ones following the pattern
@@ -270,6 +277,47 @@ default void typeConfiguration(ResolvableType resolvableType, Consumer<AotTypeCo
 	 */
 	void contributeTypeConfigurations(GenerationContext generationContext);
 
+	/**
+	 * Builder to create {@link AotContext} instances.
+	 *
+	 * @since 4.0.5
+	 */
+	interface AotContextBuilder {
+
+		/**
+		 * Configure the {@link BeanFactory} to use. Must be provided.
+		 *
+		 * @param beanFactory the bean factory to use.
+		 * @return this builder.
+		 */
+		AotContextBuilder beanFactory(BeanFactory beanFactory);
+
+		/**
+		 * Configure the {@link Environment} to use. Defaults to {@link StandardEnvironment} if not configured.
+		 *
+		 * @param environment the environment to use.
+		 * @return this builder.
+		 */
+		AotContextBuilder environment(Environment environment);
+
+		/**
+		 * Provide a consumer to customize {@link TypeCollector}.
+		 *
+		 * @param typeCollectorConsumer the consumer to configure {@link TypeCollector}.
+		 * @return this builder.
+		 */
+		AotContextBuilder customizeTypeCollector(Consumer<TypeCollector> typeCollectorConsumer);
+
+		/**
+		 * Build the {@link AotContext} to use. Building the AotContext requires {@link #beanFactory(BeanFactory)} to be
+		 * configured.
+		 *
+		 * @return the build {@link AotContext} to use.
+		 */
+		AotContext build();
+
+	}
+
 	/**
 	 * Type-based introspector to resolve {@link Class} from a type name and to introspect the bean factory for presence
 	 * of beans.

src/main/java/org/springframework/data/aot/AotMappingContext.java

@@ -36,6 +36,7 @@
 import org.springframework.data.mapping.model.EntityInstantiators;
 import org.springframework.data.mapping.model.Property;
 import org.springframework.data.mapping.model.SimpleTypeHolder;
+import org.springframework.data.util.Lazy;
 import org.springframework.data.util.TypeCollector;
 
 /**
@@ -48,17 +49,18 @@
 class AotMappingContext extends
 		AbstractMappingContext<BasicPersistentEntity<?, AotMappingContext.AotPersistentProperty>, AotMappingContext.AotPersistentProperty> {
 
+	private final static Lazy<TypeCollector> DEFAULT_TYPE_COLLECTOR = Lazy.of(TypeCollector::new);
 	private final EntityInstantiators instantiators = new EntityInstantiators();
 	private final AotAccessorFactory propertyAccessorFactory = new AotAccessorFactory();
 	private final Predicate<Class<?>> typeCollectorFilter;
 
 	AotMappingContext() {
-		this(TypeCollector.typeFilter(typeCollector -> {}));
+		this(DEFAULT_TYPE_COLLECTOR.get().getTypeFilter());
 	}
 
 	/**
-	 * @param typeFilter used for filtering during {@link #shouldCreatePersistentEntityFor(TypeInformation)}
-	 * @since 4.1
+	 * @param typeFilter used for filtering during {@link #shouldCreatePersistentEntityFor(TypeInformation)}.
+	 * @since 4.0.5
 	 */
 	AotMappingContext(Predicate<Class<?>> typeFilter) {
 		this.typeCollectorFilter = typeFilter;

src/main/java/org/springframework/data/aot/DefaultAotContextBuilder.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2026-present the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.aot;
+
+import java.util.function.Consumer;
+
+import org.jspecify.annotations.Nullable;
+
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.core.env.Environment;
+import org.springframework.core.env.StandardEnvironment;
+import org.springframework.data.util.TypeCollector;
+import org.springframework.util.Assert;
+
+/**
+ * Default implementation of {@link AotContext.AotContextBuilder}.
+ *
+ * @author Mark Paluch
+ * @since 4.0.5
+ */
+class DefaultAotContextBuilder implements AotContext.AotContextBuilder {
+
+	private @Nullable BeanFactory beanFactory;
+
+	private @Nullable Environment environment;
+
+	private Consumer<TypeCollector> typeCollectorConsumer = it -> {};
+
+	@Override
+	public AotContext.AotContextBuilder beanFactory(BeanFactory beanFactory) {
+		this.beanFactory = beanFactory;
+		return this;
+	}
+
+	@Override
+	public AotContext.AotContextBuilder environment(Environment environment) {
+		this.environment = environment;
+		return this;
+	}
+
+	@Override
+	public AotContext.AotContextBuilder customizeTypeCollector(Consumer<TypeCollector> typeCollectorConsumer) {
+		this.typeCollectorConsumer = typeCollectorConsumer;
+		return this;
+	}
+
+	@Override
+	public AotContext build() {
+
+		Assert.notNull(this.beanFactory, "BeanFactory must not be null");
+
+		Environment environment = this.environment;
+		if (environment == null) {
+			environment = new StandardEnvironment();
+		}
+
+		AotMappingContext mappingContext = new AotMappingContext(
+				TypeCollector.create(typeCollectorConsumer).getTypeFilter());
+		return new DefaultAotContext(beanFactory, environment, mappingContext);
+	}
+
+}

src/main/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessor.java

@@ -84,7 +84,7 @@ public void setEnvironment(Environment environment) {
 
 		DefaultAotContext aotContext = new DefaultAotContext(registeredBean.getBeanFactory(),
 				getConfiguredEnvironmentOrTryToResolveOne(registeredBean),
-				new AotMappingContext(TypeCollector.typeFilter(typeCollectorCustomizer())));
+				new AotMappingContext(TypeCollector.create(typeCollectorCustomizer()).getTypeFilter()));
 		return contribute(aotContext, resolveManagedTypes(registeredBean), registeredBean);
 	}
 

src/main/java/org/springframework/data/repository/config/RepositoryRegistrationAotProcessor.java

@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 import org.apache.commons.logging.Log;
@@ -51,6 +52,7 @@
 import org.springframework.data.repository.core.RepositoryInformation;
 import org.springframework.data.repository.core.support.RepositoryFactoryBeanSupport;
 import org.springframework.data.repository.core.support.RepositoryFragment;
+import org.springframework.data.util.TypeCollector;
 import org.springframework.data.util.TypeContributor;
 import org.springframework.data.util.TypeUtils;
 import org.springframework.util.Assert;
@@ -274,6 +276,16 @@ protected void configureTypeContribution(Class<?> type, AotContext aotContext) {
 		aotContext.typeConfiguration(type, config -> config.forDataBinding().contributeAccessors().forQuerydsl());
 	}
 
+	/**
+	 * Customization hook to configure {@link TypeCollector}.
+	 *
+	 * @return a {@link Consumer} to customize the {@link TypeCollector}, must not be {@literal null}.
+	 * @since 4.0
+	 */
+	protected Consumer<TypeCollector> typeCollectorCustomizer() {
+		return typeCollector -> {};
+	}
+
 	/**
 	 * This method allows for the creation to be overridden by subclasses.
 	 *
@@ -344,10 +356,13 @@ private static void registerReflectiveHints(Class<?> typeToRegister, GenerationC
 							.formatted(bean.getBeanName()));
 			return null;
 		}
+
+		AotContext aotContext = AotContext.builder().beanFactory(bean.getBeanFactory()).environment(environment)
+				.customizeTypeCollector(typeCollectorCustomizer()).build();
+
 		RepositoryInformation repositoryInformation = reader.getRepositoryInformation();
 		DefaultAotRepositoryContext repositoryContext = new DefaultAotRepositoryContext(bean, repositoryInformation,
-				extension.getModuleName(), AotContext.from(bean.getBeanFactory(), environment),
-				configuration.getConfigurationSource());
+				extension.getModuleName(), aotContext, configuration.getConfigurationSource());
 
 		repositoryContext.setIdentifyingAnnotations(extension.getIdentifyingAnnotations());