A new CodeGenerator.createFile method with more flexibility

api/src/main/kotlin/com/google/devtools/ksp/processing/CodeGenerator.kt

@@ -63,6 +63,38 @@ interface CodeGenerator {
         extensionName: String = "kt"
     ): OutputStream
 
+    /**
+     * Creates a file which is managed by [CodeGenerator]
+     *
+     * Sources of corresponding [KSNode]s which are obtained directly from [Resolver] need to be specified.
+     * Namely, the containing files of those [KSNode]s who are obtained from:
+     *   * [Resolver.getAllFiles]
+     *   * [Resolver.getSymbolsWithAnnotation]
+     *   * [Resolver.getClassDeclarationByName]
+     *
+     * Instead of requiring processors to specify all source files which are relevant in generating the given output,
+     * KSP traces dependencies automatically and only needs to know those sources that only processors know what they
+     * are for. If a [KSFile] is indirectly obtained through other [KSNode]s, it hasn't to be specified for the given
+     * output, even if its contents contribute to the generation of the output.
+     *
+     * For example, a processor generates an output `O` after reading class `A` in `A.kt` and class `B` in `B.kt`,
+     * where `A` extends `B`. The processor got `A` by [Resolver.getSymbolsWithAnnotation] and then got `B` by
+     * [KSClassDeclaration.superTypes] from `A`. Because the inclusion of `B` is due to `A`, `B.kt` needn't to be
+     * specified in [dependencies] for `O`. Note that specifying `B.kt` in this case doesn't hurt, it is only unnecessary.
+     *
+     * @param dependencies are [KSFile]s from which this output is built. Only those that are obtained directly
+     *                     from [Resolver] are required.
+     * @param path corresponds to the relative path of the generated file; includes the full file name
+     * @param fileType determines the target directory to store the file
+     * @return OutputStream for writing into files.
+     * @see [CodeGenerator] for more details.
+     */
+    fun createNewFile(
+        dependencies: Dependencies,
+        path: String,
+        fileType: FileType
+    ): OutputStream
+
     /**
      * Associate [sources] to an output file.
      *
@@ -76,6 +108,17 @@ interface CodeGenerator {
      */
     fun associate(sources: List<KSFile>, packageName: String, fileName: String, extensionName: String = "kt")
 
+    /**
+     * Associate [sources] to an output file.
+     *
+     * @param sources are [KSFile]s from which this output is built. Only those that are obtained directly
+     *                     from [Resolver] are required.
+     * @param path corresponds to the relative path of the generated file; includes the full file name
+     * @param fileType determines the target directory where the file should exist
+     * @see [CodeGenerator] for more details.
+     */
+    fun associate(sources: List<KSFile>, path: String, fileType: FileType)
+
     /**
      * Associate [classes] to an output file.
      *

api/src/main/kotlin/com/google/devtools/ksp/processing/FileType.kt

@@ -0,0 +1,12 @@
+package com.google.devtools.ksp.processing
+
+/**
+ * Used to determine where files should be stored when being
+ * created through the [CodeGenerator].
+ */
+enum class FileType {
+    CLASS,
+    JAVA_SOURCE,
+    KOTLIN_SOURCE,
+    RESOURCE
+}

common-util/src/main/kotlin/com/google/devtools/ksp/processing/impl/CodeGeneratorImpl.kt

@@ -20,6 +20,7 @@ package com.google.devtools.ksp.processing.impl
 import com.google.devtools.ksp.NoSourceFile
 import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.processing.FileType
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.KSFile
 import java.io.File
@@ -54,13 +55,25 @@ class CodeGeneratorImpl(
     fun pathOf(packageName: String, fileName: String, extensionName: String): String {
         val packageDirs = if (packageName != "") "${packageName.split(".").joinToString(separator)}$separator" else ""
         val extension = if (extensionName != "") ".$extensionName" else ""
-        val typeRoot = when (extensionName) {
-            "class" -> classDir
-            "java" -> javaDir
-            "kt" -> kotlinDir
-            else -> resourcesDir
-        }.path
-        return "$typeRoot$separator$packageDirs$fileName$extension"
+        return "$packageDirs$fileName$extension"
+    }
+
+    fun extensionToType(extensionName: String): FileType {
+        return when (extensionName) {
+            "class" -> FileType.CLASS
+            "java" -> FileType.JAVA_SOURCE
+            "kt" -> FileType.KOTLIN_SOURCE
+            else -> FileType.RESOURCE
+        }
+    }
+
+    fun baseDirOf(fileType: FileType): File {
+        return when (fileType) {
+            FileType.CLASS -> classDir
+            FileType.JAVA_SOURCE -> javaDir
+            FileType.KOTLIN_SOURCE -> kotlinDir
+            FileType.RESOURCE -> resourcesDir
+        }
     }
 
     override fun createNewFile(
@@ -69,8 +82,15 @@ class CodeGeneratorImpl(
         fileName: String,
         extensionName: String
     ): OutputStream {
-        val path = pathOf(packageName, fileName, extensionName)
-        val file = File(path)
+        return createNewFile(dependencies, pathOf(packageName, fileName, extensionName), extensionToType(extensionName))
+    }
+
+    override fun createNewFile(
+        dependencies: Dependencies,
+        path: String,
+        fileType: FileType
+    ): OutputStream {
+        val file = File(baseDirOf(fileType), path)
         if (path in fileMap) {
             throw FileAlreadyExistsException(file)
         }
@@ -89,14 +109,17 @@ class CodeGeneratorImpl(
                 dependencies.originatingFiles
             }
         }
-        associate(sources, path)
+        associate(sources, file)
         fileOutputStreamMap[path] = fileMap[path]!!.outputStream()
         return fileOutputStreamMap[path]!!
     }
 
     override fun associate(sources: List<KSFile>, packageName: String, fileName: String, extensionName: String) {
-        val path = pathOf(packageName, fileName, extensionName)
-        associate(sources, path)
+        associate(sources, pathOf(packageName, fileName, extensionName), extensionToType(extensionName))
+    }
+
+    override fun associate(sources: List<KSFile>, path: String, fileType: FileType) {
+        associate(sources, File(baseDirOf(fileType), path))
     }
 
     override fun associateWithClasses(
@@ -109,14 +132,14 @@ class CodeGeneratorImpl(
         val files = classes.map {
             it.containingFile ?: NoSourceFile(projectBase, it.qualifiedName?.asString().toString())
         }
-        associate(files, path)
+        associate(files, File(path))
     }
 
-    private fun associate(sources: List<KSFile>, outputPath: String) {
+    private fun associate(sources: List<KSFile>, outputPath: File) {
         if (!isIncremental)
             return
 
-        val output = File(outputPath).relativeTo(projectBase)
+        val output = outputPath.relativeTo(projectBase)
         sources.forEach { source ->
             sourceToOutputs.getOrPut(File(source.filePath).relativeTo(projectBase)) { mutableSetOf() }.add(output)
         }

common-util/src/test/kotlin/com/google/devtools/ksp/processing/impl/CodeGeneratorImplTest.kt

@@ -0,0 +1,96 @@
+package com.google.devtools.ksp.processing.impl
+
+import com.google.devtools.ksp.AnyChanges
+import com.google.devtools.ksp.processing.CodeGenerator
+import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.processing.FileType
+import org.junit.Assert
+import org.junit.Before
+import org.junit.Test
+import java.io.File
+import java.nio.file.Files
+
+class CodeGeneratorImplTest {
+
+    lateinit var codeGenerator: CodeGenerator
+    lateinit var baseDir: File
+
+    @Before
+    fun setup() {
+        baseDir = Files.createTempDirectory("project").toFile()
+        val classesDir = File(baseDir, "classes")
+        classesDir.mkdir()
+        val javaDir = File(baseDir, "java")
+        javaDir.mkdir()
+        val kotlinDir = File(baseDir, "kotlin")
+        kotlinDir.mkdir()
+        val resourcesDir = File(baseDir, "resources")
+        resourcesDir.mkdir()
+        codeGenerator = CodeGeneratorImpl(
+            classesDir,
+            javaDir,
+            kotlinDir,
+            resourcesDir,
+            baseDir,
+            AnyChanges(baseDir),
+            emptyList(),
+            true
+        )
+    }
+
+    @Test
+    fun testCreatingAFile() {
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "java")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "kt")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "class")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a.b.c", "Test", "")
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/Test"), files[3])
+    }
+
+    @Test
+    fun testCreatingAFileWithSlash() {
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "java")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "kt")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "class")
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c", "Test", "")
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/Test"), files[3])
+    }
+
+    @Test
+    fun testCreatingAFileWithPath() {
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/Test.java", FileType.JAVA_SOURCE)
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/Test.kt", FileType.KOTLIN_SOURCE)
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/Test.class", FileType.CLASS)
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/Test", FileType.RESOURCE)
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/Test"), files[3])
+    }
+
+    @Test
+    fun testCreatingAFileWithPathAndDots() {
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test.java", FileType.JAVA_SOURCE)
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test.kt", FileType.KOTLIN_SOURCE)
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test.class", FileType.CLASS)
+        codeGenerator.createNewFile(Dependencies.ALL_FILES, "a/b/c/dir.with.dot/Test", FileType.RESOURCE)
+
+        val files = codeGenerator.generatedFile.toList()
+        Assert.assertEquals(File(baseDir, "java/a/b/c/dir.with.dot/Test.java"), files[0])
+        Assert.assertEquals(File(baseDir, "kotlin/a/b/c/dir.with.dot/Test.kt"), files[1])
+        Assert.assertEquals(File(baseDir, "classes/a/b/c/dir.with.dot/Test.class"), files[2])
+        Assert.assertEquals(File(baseDir, "resources/a/b/c/dir.with.dot/Test"), files[3])
+    }
+}