Let fluent setters of SimpleStepBuilder return proper type

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/AbstractTaskletStepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2018 the original author or authors.
+ * Copyright 2012-2022 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.
@@ -51,8 +51,7 @@
  * @since 2.2
  * @param <B> the type of builder represented
  */
-public abstract class AbstractTaskletStepBuilder<B extends AbstractTaskletStepBuilder<B>>
-		extends StepBuilderHelper<AbstractTaskletStepBuilder<B>> {
+public abstract class AbstractTaskletStepBuilder<B extends AbstractTaskletStepBuilder<B>> extends StepBuilderHelper<B> {
 
 	protected Set<ChunkListener> chunkListeners = new LinkedHashSet<>();
 
@@ -137,9 +136,9 @@ protected void registerStepListenerAsChunkListener() {
 	 * @param listener the listener to register
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> listener(ChunkListener listener) {
+	public B listener(ChunkListener listener) {
 		chunkListeners.add(listener);
-		return this;
+		return self();
 	}
 
 	/**
@@ -162,19 +161,17 @@ public B listener(Object listener) {
 			this.listener((ChunkListener) factory.getObject());
 		}
 
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
 	/**
 	 * Register a stream for callbacks that manage restart data.
 	 * @param stream the stream to register
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> stream(ItemStream stream) {
+	public B stream(ItemStream stream) {
 		streams.add(stream);
-		return this;
+		return self();
 	}
 
 	/**
@@ -183,9 +180,9 @@ public AbstractTaskletStepBuilder<B> stream(ItemStream stream) {
 	 * @param taskExecutor the task executor to register
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> taskExecutor(TaskExecutor taskExecutor) {
+	public B taskExecutor(TaskExecutor taskExecutor) {
 		this.taskExecutor = taskExecutor;
-		return this;
+		return self();
 	}
 
 	/**
@@ -196,9 +193,9 @@ public AbstractTaskletStepBuilder<B> taskExecutor(TaskExecutor taskExecutor) {
 	 * @param throttleLimit maximum number of concurrent tasklet executions allowed
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> throttleLimit(int throttleLimit) {
+	public B throttleLimit(int throttleLimit) {
 		this.throttleLimit = throttleLimit;
-		return this;
+		return self();
 	}
 
 	/**
@@ -207,9 +204,9 @@ public AbstractTaskletStepBuilder<B> throttleLimit(int throttleLimit) {
 	 * @param exceptionHandler the exception handler
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> exceptionHandler(ExceptionHandler exceptionHandler) {
+	public B exceptionHandler(ExceptionHandler exceptionHandler) {
 		this.exceptionHandler = exceptionHandler;
-		return this;
+		return self();
 	}
 
 	/**
@@ -218,9 +215,9 @@ public AbstractTaskletStepBuilder<B> exceptionHandler(ExceptionHandler exception
 	 * @param repeatTemplate a repeat template with rules for iterating
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> stepOperations(RepeatOperations repeatTemplate) {
+	public B stepOperations(RepeatOperations repeatTemplate) {
 		this.stepOperations = repeatTemplate;
-		return this;
+		return self();
 	}
 
 	/**
@@ -230,9 +227,9 @@ public AbstractTaskletStepBuilder<B> stepOperations(RepeatOperations repeatTempl
 	 * @param transactionAttribute a transaction attribute set
 	 * @return this for fluent chaining
 	 */
-	public AbstractTaskletStepBuilder<B> transactionAttribute(TransactionAttribute transactionAttribute) {
+	public B transactionAttribute(TransactionAttribute transactionAttribute) {
 		this.transactionAttribute = transactionAttribute;
-		return this;
+		return self();
 	}
 
 	/**

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/FaultTolerantStepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2021 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -205,9 +205,7 @@ public SimpleStepBuilder<I, O> listener(Object listener) {
 			skipListeners.add((SkipListener) factory.getObject());
 		}
 
-		@SuppressWarnings("unchecked")
-		SimpleStepBuilder<I, O> result = this;
-		return result;
+		return this;
 	}
 
 	/**
@@ -227,8 +225,7 @@ public FaultTolerantStepBuilder<I, O> listener(ChunkListener listener) {
 	}
 
 	@Override
-	public AbstractTaskletStepBuilder<SimpleStepBuilder<I, O>> transactionAttribute(
-			TransactionAttribute transactionAttribute) {
+	public SimpleStepBuilder<I, O> transactionAttribute(TransactionAttribute transactionAttribute) {
 		return super.transactionAttribute(getTransactionAttribute(transactionAttribute));
 	}
 
@@ -397,7 +394,7 @@ public FaultTolerantStepBuilder<I, O> processorNonTransactional() {
 	}
 
 	@Override
-	public AbstractTaskletStepBuilder<SimpleStepBuilder<I, O>> stream(ItemStream stream) {
+	public SimpleStepBuilder<I, O> stream(ItemStream stream) {
 		if (stream instanceof ItemReader<?>) {
 			if (!streamIsReader) {
 				streamIsReader = true;

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/FlowStepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2011 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -69,4 +69,9 @@ public Step build() {
 		return step;
 	}
 
+	@Override
+	protected FlowStepBuilder self() {
+		return this;
+	}
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/JobStepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2011 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -114,4 +114,9 @@ public Step build() {
 
 	}
 
+	@Override
+	protected JobStepBuilder self() {
+		return this;
+	}
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/PartitionStepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2021 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -222,6 +222,11 @@ public Step build() {
 
 	}
 
+	@Override
+	protected PartitionStepBuilder self() {
+		return this;
+	}
+
 	protected TaskExecutor getTaskExecutor() {
 		return taskExecutor;
 	}

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/SimpleStepBuilder.java

@@ -246,7 +246,6 @@ public SimpleStepBuilder<I, O> readerIsTransactionalQueue() {
 	 * @param listener the object that has a method configured with listener annotation
 	 * @return this for fluent chaining
 	 */
-	@SuppressWarnings("unchecked")
 	@Override
 	public SimpleStepBuilder<I, O> listener(Object listener) {
 		super.listener(listener);
@@ -268,9 +267,7 @@ public SimpleStepBuilder<I, O> listener(Object listener) {
 			itemListeners.add((StepListener) factory.getObject());
 		}
 
-		@SuppressWarnings("unchecked")
-		SimpleStepBuilder<I, O> result = this;
-		return result;
+		return this;
 	}
 
 	/**
@@ -315,6 +312,11 @@ public SimpleStepBuilder<I, O> chunkOperations(RepeatOperations repeatTemplate)
 		return this;
 	}
 
+	@Override
+	protected SimpleStepBuilder<I, O> self() {
+		return this;
+	}
+
 	protected RepeatOperations createChunkOperations() {
 		RepeatOperations repeatOperations = chunkOperations;
 		if (repeatOperations == null) {

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/StepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2013 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -124,4 +124,9 @@ public FlowStepBuilder flow(Flow flow) {
 		return new FlowStepBuilder(this).flow(flow);
 	}
 
+	@Override
+	protected StepBuilder self() {
+		return this;
+	}
+
 }

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/StepBuilderHelper.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2014 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -64,23 +64,17 @@ protected StepBuilderHelper(StepBuilderHelper<?> parent) {
 
 	public B repository(JobRepository jobRepository) {
 		properties.jobRepository = jobRepository;
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
 	public B transactionManager(PlatformTransactionManager transactionManager) {
 		properties.transactionManager = transactionManager;
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
 	public B startLimit(int startLimit) {
 		properties.startLimit = startLimit;
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
 	/**
@@ -99,25 +93,21 @@ public B listener(Object listener) {
 			properties.addStepExecutionListener((StepExecutionListener) factory.getObject());
 		}
 
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
 	public B listener(StepExecutionListener listener) {
 		properties.addStepExecutionListener(listener);
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
 	public B allowStartIfComplete(boolean allowStartIfComplete) {
 		properties.allowStartIfComplete = allowStartIfComplete;
-		@SuppressWarnings("unchecked")
-		B result = (B) this;
-		return result;
+		return self();
 	}
 
+	protected abstract B self();
+
 	protected String getName() {
 		return properties.name;
 	}

spring-batch-core/src/main/java/org/springframework/batch/core/step/builder/TaskletStepBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2006-2011 the original author or authors.
+ * Copyright 2006-2022 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.
@@ -45,6 +45,11 @@ public TaskletStepBuilder tasklet(Tasklet tasklet) {
 		return this;
 	}
 
+	@Override
+	protected TaskletStepBuilder self() {
+		return this;
+	}
+
 	@Override
 	protected Tasklet createTasklet() {
 		return tasklet;