diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/FunctionCodegen.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/FunctionCodegen.kt new file mode 100644 index 0000000000..a2605de2ad --- /dev/null +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/FunctionCodegen.kt @@ -0,0 +1,128 @@ +package org.jetbrains.kotlinx.dataframe.codeGen + +import com.squareup.kotlinpoet.ClassName +import com.squareup.kotlinpoet.FunSpec +import com.squareup.kotlinpoet.ParameterSpec +import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy +import com.squareup.kotlinpoet.TypeName +import com.squareup.kotlinpoet.TypeSpec +import com.squareup.kotlinpoet.asTypeName +import com.squareup.kotlinpoet.typeNameOf +import org.jetbrains.dataframe.impl.codeGen.CodeGenerator +import org.jetbrains.kotlinx.dataframe.DataFrame +import org.jetbrains.kotlinx.dataframe.api.schema +import org.jetbrains.kotlinx.dataframe.impl.DELIMITERS_REGEX +import org.jetbrains.kotlinx.dataframe.impl.codeGen.toStandaloneSnippet +import org.jetbrains.kotlinx.dataframe.impl.joinToCamelCaseString +import org.jetbrains.kotlinx.dataframe.impl.schema.intersectSchemas +import java.io.File +import kotlin.reflect.KClass +import kotlin.reflect.KFunction +import kotlin.reflect.jvm.javaMethod + +public inline fun generateCode(noinline f: () -> DataFrame, sourceSet: File, name: String) { + return generateCode(f, name, T::class, sourceSet) +} + +public inline fun generateCode( + noinline f: (I1) -> DataFrame, sourceSet: File, name: String, arguments: List = emptyList() +) { + return generateCode(f, name, arguments, typeNameOf(), T::class, sourceSet) +} + +public fun generateCode(f: Function2>, arguments: List>) { + +} + +@PublishedApi +internal fun generateCode( + f: (I1) -> DataFrame, + name: String, + arguments: List = emptyList(), + inputName: TypeName, + marker: KClass<*>, + sourceSet: File +) { + val schema = if (arguments.isNotEmpty()) { + arguments.map { f(it).schema() }.intersectSchemas() + } else { + MarkersExtractor.get(marker).schema + } + val packageName = (f as KFunction<*>).javaMethod?.declaringClass?.packageName!! + val functionName = (f as KFunction<*>).javaMethod?.name!! + val method = object : DefaultReadDfMethod { + override val additionalImports: List = emptyList() + + override fun toDeclaration(marker: Marker, visibility: String): String { + val cl = (f as KFunction<*>).javaMethod?.declaringClass?.name!! + val format = + """ + val f = this::class.java.classLoader.loadClass("$cl").getDeclaredMethod("$functionName", ${inputName}::class.java) + f.isAccessible = true + return (f.invoke(null, v) as DataFrame<*>).cast() + """.trimIndent() + + val typeSpec = TypeSpec.companionObjectBuilder() + .addFunction( + FunSpec.builder("from") + .addParameter( + ParameterSpec.builder("v", inputName).build() + ) + .returns(DataFrame::class.asTypeName().parameterizedBy(ClassName(packageName, marker.shortName))) + .addCode(format, "v") + .build() + ).build() + return typeSpec.toString() + } + } + + val generator = CodeGenerator.create(useFqNames = true).generate(schema, name, true, false, true, visibility = MarkerVisibility.EXPLICIT_PUBLIC,readDfMethod = method) + + println(packageName) + println(generator.code.toStandaloneSnippet(packageName, emptyList())) + val destination = sourceSet + val targetPackage = File(destination, packageName.replace(".", "/")) + val file = File(targetPackage, "${name.split(DELIMITERS_REGEX).joinToString("")}.Generated.kt") + file.writeText(generator.code.toStandaloneSnippet(packageName, emptyList())) +} + +@PublishedApi +internal fun generateCode(f: () -> DataFrame, name: String, marker: KClass<*>, sourceSet: File) { + // refine nullability from compile time schema +// val schema = MarkersExtractor.get(marker).schema + val schema = f().schema() + val packageName = (f as KFunction<*>).javaMethod?.declaringClass?.packageName!! + val functionName = (f as KFunction<*>).javaMethod?.name!! + val method = object : DefaultReadDfMethod { + override val additionalImports: List = emptyList() + + override fun toDeclaration(marker: Marker, visibility: String): String { + val cl = (f as KFunction<*>).javaMethod?.declaringClass?.name!! + val format = + """ + val f = this::class.java.classLoader.loadClass("$cl").getDeclaredMethod("$functionName") + f.isAccessible = true + return (f.invoke(null) as DataFrame<*>).cast() + """.trimIndent() + + val typeSpec = TypeSpec.companionObjectBuilder() + .addFunction( + FunSpec.builder("sample") + .returns(DataFrame::class.asTypeName().parameterizedBy(ClassName(packageName, marker.shortName))) + .addCode(format, "v") + .build() + ).build() + return typeSpec.toString() + } + } + + val generator = CodeGenerator.create(useFqNames = true).generate(schema, name, true, false, true, visibility = MarkerVisibility.EXPLICIT_PUBLIC,readDfMethod = method) + + println(packageName) + println(generator.code.toStandaloneSnippet(packageName, emptyList())) + val destination = sourceSet + val targetPackage = File(destination, packageName.replace(".", "/")) + val file = File(targetPackage, "${name.split(DELIMITERS_REGEX).joinToString("")}.Generated.kt") + file.writeText(generator.code.toStandaloneSnippet(packageName, emptyList())) +} +