Add configurable Predicate for type inclusion to TypeCollector.

Author: mp911de

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

@@ -18,6 +18,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Set;
+import java.util.function.Consumer;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -35,6 +36,7 @@
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.data.domain.ManagedTypes;
 import org.springframework.data.util.Lazy;
+import org.springframework.data.util.TypeCollector;
 import org.springframework.data.util.TypeContributor;
 import org.springframework.data.util.TypeUtils;
 import org.springframework.util.ClassUtils;
@@ -125,9 +127,18 @@ private ManagedTypes resolveManagedTypes(RegisteredBean registeredBean) {
 	 */
 	protected BeanRegistrationAotContribution contribute(AotContext aotContext, ManagedTypes managedTypes,
 			RegisteredBean registeredBean) {
-		return new ManagedTypesRegistrationAotContribution(aotContext, managedTypes, registeredBean, this::contributeType);
+		return new ManagedTypesRegistrationAotContribution(aotContext, managedTypes, registeredBean,
+				typeCollectorCustomizer(), this::contributeType);
 	}
 
+	/**
+	 * Customization hook to configure {@link TypeCollector}.
+	 *
+	 * @return a {@link Consumer} to customize the {@link TypeCollector}, must not be {@literal null}.
+	 */
+	protected Consumer<TypeCollector> typeCollectorCustomizer() {
+		return typeCollector -> {};
+	}
 	/**
 	 * Hook to contribute configuration for a given {@literal type}.
 	 *
@@ -142,14 +153,26 @@ protected void contributeType(ResolvableType type, GenerationContext generationC
 
 		Set<String> annotationNamespaces = Collections.singleton(TypeContributor.DATA_NAMESPACE);
 
-		aotContext.typeConfiguration(type, config -> config.forDataBinding() //
-				.contributeAccessors() //
-				.forQuerydsl().contribute(environment.get(), generationContext));
+		configureTypeContribution(type.toClass(), aotContext);
+
+		aotContext.typeConfiguration(type, config -> {
+			config.contribute(environment.get(), generationContext);
+		});
 
 		TypeUtils.resolveUsedAnnotations(type.toClass()).forEach(
 				annotation -> TypeContributor.contribute(annotation.getType(), annotationNamespaces, generationContext));
 	}
 
+	/**
+	 * Customization hook to configure the {@link TypeContributor} used to register the given {@literal type}.
+	 *
+	 * @param type the class to configure the contribution for.
+	 * @param aotContext AOT context for type configuration.
+	 */
+	protected void configureTypeContribution(Class<?> type, AotContext aotContext) {
+		aotContext.typeConfiguration(type, config -> config.forDataBinding().contributeAccessors().forQuerydsl());
+	}
+
 	protected boolean isMatch(@Nullable Class<?> beanType, @Nullable String beanName) {
 		return matchesByType(beanType) && matchesPrefix(beanName);
 	}

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

@@ -18,11 +18,12 @@
 import java.lang.reflect.Method;
 import java.util.Collections;
 import java.util.List;
-import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 
 import javax.lang.model.element.Modifier;
 
 import org.jspecify.annotations.Nullable;
+
 import org.springframework.aot.generate.AccessControl;
 import org.springframework.aot.generate.GeneratedMethod;
 import org.springframework.aot.generate.GenerationContext;
@@ -76,15 +77,18 @@ class ManagedTypesRegistrationAotContribution implements RegisteredBeanAotContri
 	private final AotContext aotContext;
 	private final ManagedTypes managedTypes;
 	private final Lazy<List<Class<?>>> sourceTypes;
+	private final Consumer<TypeCollector> typeCollectorCustomizer;
 	private final TypeRegistration contributionAction;
 	private final RegisteredBean source;
 
 	public ManagedTypesRegistrationAotContribution(AotContext aotContext, ManagedTypes managedTypes,
-			RegisteredBean registeredBean, TypeRegistration contributionAction) {
+			RegisteredBean registeredBean, Consumer<TypeCollector> typeCollectorCustomizer,
+			TypeRegistration contributionAction) {
 
 		this.aotContext = aotContext;
 		this.managedTypes = managedTypes;
 		this.sourceTypes = Lazy.of(managedTypes::toList);
+		this.typeCollectorCustomizer = typeCollectorCustomizer;
 		this.contributionAction = contributionAction;
 		this.source = registeredBean;
 	}
@@ -95,18 +99,15 @@ public void applyTo(GenerationContext generationContext, BeanRegistrationCode be
 		List<Class<?>> types = sourceTypes.get();
 
 		if (!types.isEmpty()) {
-			TypeCollector.inspect(types).forEach(type -> contributionAction.register(type, generationContext, aotContext));
+			TypeCollector.inspect(typeCollectorCustomizer, types)
+					.forEach(type -> contributionAction.register(type, generationContext, aotContext));
 		}
 	}
 
 	@Override
 	public BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(GenerationContext generationContext,
 			BeanRegistrationCodeFragments codeFragments) {
 
-		if (managedTypes == null) {
-			return codeFragments;
-		}
-
 		ManagedTypesInstanceCodeFragment fragment = new ManagedTypesInstanceCodeFragment(sourceTypes.get(), source,
 				codeFragments);
 		return fragment.canGenerateCode() ? fragment : codeFragments;

src/main/java/org/springframework/data/util/TypeCollector.java

@@ -94,8 +94,40 @@ public static ReachableTypes inspect(Class<?>... types) {
 		return inspect(Arrays.asList(types));
 	}
 
+	/**
+	 * Inspect the given type and resolve those reachable via fields, methods, generics, ...
+	 *
+	 * @param types the types to inspect
+	 * @return a type model collector for the type
+	 */
 	public static ReachableTypes inspect(Collection<Class<?>> types) {
-		return new ReachableTypes(new TypeCollector(), types);
+		return inspect(it -> {}, types);
+	}
+
+	/**
+	 * Inspect the given type and resolve those reachable via fields, methods, generics, ...
+	 *
+	 * @param collectorCustomizer the customizer function to configure the {@link TypeCollector}.
+	 * @param types the types to inspect.
+	 * @return a type model collector for the type.
+	 * @since 4.0
+	 */
+	public static ReachableTypes inspect(Consumer<TypeCollector> collectorCustomizer, Class<?>... types) {
+		return inspect(collectorCustomizer, Arrays.asList(types));
+	}
+
+	/**
+	 * Inspect the given type and resolve those reachable via fields, methods, generics, ...
+	 *
+	 * @param collectorCustomizer the customizer function to configure the {@link TypeCollector}.
+	 * @param types the types to inspect.
+	 * @return a type model collector for the type.
+	 * @since 4.0
+	 */
+	public static ReachableTypes inspect(Consumer<TypeCollector> collectorCustomizer, Collection<Class<?>> types) {
+		TypeCollector typeCollector = new TypeCollector();
+		collectorCustomizer.accept(typeCollector);
+		return new ReachableTypes(typeCollector, types);
 	}
 
 	private void process(Class<?> root, Consumer<ResolvableType> consumer) {
@@ -225,30 +257,45 @@ private Predicate<Field> createFieldFilter() {
 		return (Predicate) excludedFieldPredicate.negate();
 	}
 
+	/**
+	 * Container for reachable types starting from a set of root types.
+	 */
 	public static class ReachableTypes {
 
 		private final Iterable<Class<?>> roots;
 		private final Lazy<List<Class<?>>> reachableTypes = Lazy.of(this::collect);
 		private final TypeCollector typeCollector;
 
-		public ReachableTypes(TypeCollector typeCollector, Iterable<Class<?>> roots) {
+		ReachableTypes(TypeCollector typeCollector, Iterable<Class<?>> roots) {
 
 			this.typeCollector = typeCollector;
 			this.roots = roots;
 		}
 
-		public void forEach(Consumer<ResolvableType> consumer) {
-			roots.forEach(it -> typeCollector.process(it, consumer));
+		/**
+		 * Performs the given action for each element of the reachable types until all elements have been processed or the
+		 * action throws an exception. Actions are performed in the order of iteration, if that order is specified.
+		 * Exceptions thrown by the action are relayed to the caller.
+		 *
+		 * @param action The action to be performed for each element
+		 */
+		public void forEach(Consumer<ResolvableType> action) {
+			roots.forEach(it -> typeCollector.process(it, action));
 		}
 
+		/**
+		 * Return all reachable types as list of {@link Class classes}. The resulting list is unmodifiable.
+		 *
+		 * @return an unmodifiable list of reachable types.
+		 */
 		public List<Class<?>> list() {
 			return reachableTypes.get();
 		}
 
 		private List<Class<?>> collect() {
 			List<Class<?>> target = new ArrayList<>();
 			forEach(it -> target.add(it.toClass()));
-			return target;
+			return List.copyOf(target);
 		}
 	}
 

src/test/java/org/springframework/data/aot/TypeCollectorUnitTests.java

@@ -22,7 +22,10 @@
 import org.springframework.data.util.TypeCollector;
 
 /**
+ * Unit tests for {@link TypeCollector}.
+ *
  * @author Christoph Strobl
+ * @author Mark Paluch
  */
 public class TypeCollectorUnitTests {
 
@@ -66,4 +69,12 @@ void skipsCoreFrameworkType() {
 		assertThat(TypeCollector.inspect(org.springframework.core.AliasRegistry.class).list()).isEmpty();
 	}
 
+	@Test // GH-3362
+	void appliesFilterPredicate() {
+		assertThat(TypeCollector
+				.inspect(it -> it.filterTypes(cls -> cls == EmptyType1.class || cls == TypesInMethodSignatures.class),
+						TypesInMethodSignatures.class)
+				.list()).containsOnly(TypesInMethodSignatures.class, EmptyType1.class);
+	}
+
 }