Add subfolder of task output to main sources

I am using a task to generate code into layout.buildDirectory.dir("swagger-code-gen")
Inside this directory a task will then create the subdirectories src/main/java and java files

I could depend on the tasks output using the following line:
sourceSets[SourceSet.MAIN_SOURCE_SET_NAME].java.srcDir(files(openApiGenerateTask))

Is there a way to say, that I only want to add the src/main/java subdir of the tasks output?

My current workaround is to use a sync task and extract the src/main/java dir:

val openApiSrcTask by tasks.registering(Sync::class) {
    from(openApiGenerateTask)
    include("src/main/java/**")
    eachFile { relativePath = RelativePath(true, *relativePath.segments.drop(3).toTypedArray()) }
    includeEmptyDirs = false
    into(layout.buildDirectory.dir("swagger-code-gen-java"))
}

and the use its output:
sourceSets[SourceSet.MAIN_SOURCE_SET_NAME].java.srcDir(files(openApiSrcTask))