generate refs class

kotlin-orm · Jul 16, 2023 · a783546 · a783546
1 parent c2968e8
commit a783546
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/KtormProcessorProvider.kt b/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/KtormProcessorProvider.kt
@@ -95,6 +95,8 @@ public class KtormProcessorProvider : SymbolProcessorProvider {
                 .replace(Regex(""",\s*"""), ", ")
                 .replace(Regex(""",\s*\)"""), ")")
                 .replace(Regex("""\s*get\(\)\s="""), " get() =")
+                .replace(Regex("""\s+=\s+"""), " = ")
+                .replace("import org.ktorm.ksp.`annotation`", "import org.ktorm.ksp.annotation")
 
             return ktLintRuleEngine.format(code)
         } catch (e: Exception) {

diff --git a/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/generator/RefsClassGenerator.kt b/ktorm-ksp-compiler/src/main/kotlin/org/ktorm/ksp/compiler/generator/RefsClassGenerator.kt
@@ -2,6 +2,7 @@ package org.ktorm.ksp.compiler.generator
 
 import com.squareup.kotlinpoet.*
 import org.ktorm.ksp.spi.TableMetadata
+import org.ktorm.schema.Table
 
 /**
  * Created by vince at Jul 15, 2023.
@@ -27,6 +28,17 @@ internal object RefsClassGenerator {
             typeSpec.addProperty(propertySpec)
         }
 
+        val funSpec = FunSpec.builder("toList")
+            .addKdoc("Return all referenced tables as a list.")
+            .returns(typeNameOf<List<Table<*>>>())
+            .addCode("return·listOf(")
+
+        for (column in table.columns) {
+            val refTablePropertyName = column.refTablePropertyName ?: continue
+            funSpec.addCode("%N, ", refTablePropertyName)
+        }
+
+        typeSpec.addFunction(funSpec.addCode(")").build())
         return typeSpec.build()
     }
 }