Merge pull request #247 from ermolenkodev/140

Enable multi-round DataSchema declaration processing (#140)

Author: ermolenkodev

data/jetbrains_repositories.json

@@ -0,0 +1,11 @@
+{
+  "repositories": [
+    {
+      "full_name": "JetBrains/JPS",
+      "html_url": "",
+      "stargazers_count": 10,
+      "topics": "[jetbrains, jetbrains-youtrack, youtrack, youtrack-api]",
+      "watchers": 10
+    }
+  ]
+}

plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/DataFrameSymbolProcessor.kt

@@ -13,8 +13,8 @@ class DataFrameSymbolProcessor(
 
     override fun process(resolver: Resolver): List<KSAnnotated> {
         val extensionsGenerator = ExtensionsGenerator(resolver, codeGenerator, logger)
-        val dataSchemas = extensionsGenerator.resolveValidDataSchemaDeclarations()
-        dataSchemas.forEach {
+        val (validDataSchemas, invalidDataSchemas) = extensionsGenerator.resolveDataSchemaDeclarations()
+        validDataSchemas.forEach {
             val file = it.origin.containingFile ?: return@forEach
             extensionsGenerator.generateExtensions(file, it.origin, it.properties)
         }
@@ -25,6 +25,8 @@ class DataFrameSymbolProcessor(
             dataSchemaGenerator.generateDataSchema(importStatement)
         }
 
-        return emptyList()
+        // by returning invalidDataSchemas we defer the processing of incomplete DataSchema declarations
+        // for example when DataSchema declaration references another one generated by @file:ImportDataSchema
+        return invalidDataSchemas
     }
 }

plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/ExtensionsGenerator.kt

@@ -30,19 +30,19 @@ class ExtensionsGenerator(
         val EXPECTED_VISIBILITIES = setOf(Visibility.PUBLIC, Visibility.INTERNAL)
     }
 
-    fun resolveValidDataSchemaDeclarations(): Sequence<DataSchemaDeclaration> {
+    fun resolveDataSchemaDeclarations(): Pair<Sequence<DataSchemaDeclaration>, List<KSClassDeclaration>> {
         val dataSchemaAnnotation = resolver.getKSNameFromString(DataFrameNames.DATA_SCHEMA)
         val symbols = resolver.getSymbolsWithAnnotation(dataSchemaAnnotation.asString())
 
-        return symbols
+        val (validDeclarations, invalidDeclarations) = symbols
             .filterIsInstance<KSClassDeclaration>()
-            .mapNotNull {
-                if (it.validate()) {
-                    it.toDataSchemaDeclarationOrNull()
-                } else {
-                    null
-                }
-            }
+            .partition { it.validate() }
+
+        val preprocessedDeclarations = validDeclarations
+            .asSequence()
+            .mapNotNull { it.toDataSchemaDeclarationOrNull() }
+
+        return Pair(preprocessedDeclarations, invalidDeclarations)
     }
 
     class DataSchemaDeclaration(

plugins/symbol-processor/src/test/kotlin/org/jetbrains/dataframe/ksp/DataFrameSymbolProcessorTest.kt

@@ -210,6 +210,42 @@ class DataFrameSymbolProcessorTest {
         result.successfulCompilation shouldBe true
     }
 
+    @Test
+    fun `multi-round data schema generation`() {
+        useHostedFile(jetbrainsCsv) {
+            val result = KspCompilationTestRunner.compile(
+                TestCompilationParameters(
+                    sources = listOf(
+                        SourceFile.kotlin(
+                            "MySources.kt",
+                            """
+                @file:ImportDataSchema(name = "Repo", "$it")
+
+                import org.jetbrains.kotlinx.dataframe.DataFrame
+                import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+                import org.jetbrains.kotlinx.dataframe.annotations.ImportDataSchema
+                import org.jetbrains.kotlinx.dataframe.api.print
+                import org.jetbrains.kotlinx.dataframe.io.readJson
+
+                @DataSchema
+                interface Repos {
+                    val repositories: DataFrame<Repo>
+                }
+
+                fun main() {
+                    val df: DataFrame<Repos> = DataFrame.readJson("data/jetbrains_repositories.json") as DataFrame<Repos>
+                    df.repositories[0].print()
+                }
+
+                        """.trimIndent()
+                        )
+                    )
+                )
+            )
+            result.successfulCompilation shouldBe true
+        }
+    }
+
     @Test
     fun `functional type`() {
         val result = KspCompilationTestRunner.compile(