Skip to content

Commit 4c4c251

Browse files
JolanrensenJolanrensen
committed
updated kdocs as well
1 parent 7d73d6c commit 4c4c251

File tree

3 files changed

+140
-61
lines changed

3 files changed

+140
-61
lines changed

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/Integration.kt

Lines changed: 98 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,96 @@ internal class Integration(
4545

4646
val version = options["v"]
4747

48+
private fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, argument: String): VariableName? {
49+
val code = codeWithConverter.with(argument)
50+
return if (code.isNotBlank()) {
51+
val result = execute(code)
52+
if (codeWithConverter.hasConverter) {
53+
result.name
54+
} else null
55+
} else null
56+
}
57+
58+
private fun KotlinKernelHost.execute(
59+
codeWithConverter: CodeWithConverter,
60+
property: KProperty<*>,
61+
type: String,
62+
): VariableName? {
63+
val variableName = "(${property.name}${if (property.returnType.isMarkedNullable) "!!" else ""} as $type)"
64+
return execute(codeWithConverter, variableName)
65+
}
66+
67+
private fun KotlinKernelHost.updateImportDataSchemaVariable(
68+
importDataSchema: ImportDataSchema,
69+
property: KProperty<*>,
70+
): VariableName? {
71+
val formats = supportedFormats.filterIsInstance<SupportedCodeGenerationFormat>()
72+
val name = property.name + "DataSchema"
73+
return when (
74+
val codeGenResult = CodeGenerator.urlCodeGenReader(importDataSchema.url, name, formats, true)
75+
) {
76+
is CodeGenerationReadResult.Success -> {
77+
val readDfMethod = codeGenResult.getReadDfMethod(importDataSchema.url.toExternalForm())
78+
val code = readDfMethod.additionalImports.joinToString("\n") +
79+
"\n" +
80+
codeGenResult.code
81+
82+
execute(code)
83+
execute("""DISPLAY("Data schema successfully imported as ${property.name}: $name")""")
84+
85+
name
86+
}
87+
88+
is CodeGenerationReadResult.Error -> {
89+
execute("""DISPLAY("Failed to read data schema from ${importDataSchema.url}: ${codeGenResult.reason}")""")
90+
null
91+
}
92+
}
93+
}
94+
95+
private fun KotlinKernelHost.updateAnyFrameVariable(
96+
df: AnyFrame,
97+
property: KProperty<*>,
98+
codeGen: ReplCodeGenerator,
99+
): VariableName? = execute(
100+
codeWithConverter = codeGen.process(df, property),
101+
property = property,
102+
type = "AnyFrame",
103+
)
104+
105+
private fun KotlinKernelHost.updateAnyRowVariable(
106+
row: AnyRow,
107+
property: KProperty<*>,
108+
codeGen: ReplCodeGenerator,
109+
): VariableName? = execute(
110+
codeWithConverter = codeGen.process(row, property),
111+
property = property,
112+
type = "AnyRow",
113+
)
114+
115+
private fun KotlinKernelHost.updateColumnGroupVariable(
116+
col: ColumnGroup<*>,
117+
property: KProperty<*>,
118+
codeGen: ReplCodeGenerator,
119+
): VariableName? = execute(
120+
codeWithConverter = codeGen.process(col.asDataFrame(), property),
121+
property = property,
122+
type = "ColumnGroup<*>",
123+
)
124+
125+
private fun KotlinKernelHost.updateAnyColVariable(
126+
col: AnyCol,
127+
property: KProperty<*>,
128+
codeGen: ReplCodeGenerator,
129+
): VariableName? = if (col.isColumnGroup()) {
130+
val codeWithConverter = codeGen.process(col.asColumnGroup().asDataFrame(), property).let { c ->
131+
CodeWithConverter(c.declarations) { c.converter("$it.asColumnGroup()") }
132+
}
133+
execute(codeWithConverter = codeWithConverter, property = property, type = "AnyCol")
134+
} else {
135+
null
136+
}
137+
48138
override fun Builder.onLoaded() {
49139
if (version != null) {
50140
dependencies(
@@ -152,65 +242,17 @@ internal class Integration(
152242
import("org.jetbrains.kotlinx.dataframe.dataTypes.*")
153243
import("org.jetbrains.kotlinx.dataframe.impl.codeGen.urlCodeGenReader")
154244

155-
fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, argument: String): VariableName? {
156-
val code = codeWithConverter.with(argument)
157-
return if (code.isNotBlank()) {
158-
val result = execute(code)
159-
if (codeWithConverter.hasConverter) {
160-
result.name
161-
} else null
162-
} else null
163-
}
164-
165-
fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, property: KProperty<*>): VariableName? {
166-
val variableName = property.name + if (property.returnType.isMarkedNullable) "!!" else ""
167-
return execute(codeWithConverter, variableName)
168-
}
169-
170-
updateVariable<ImportDataSchema> { importDataSchema, property ->
171-
val formats = supportedFormats.filterIsInstance<SupportedCodeGenerationFormat>()
172-
val name = property.name + "DataSchema"
173-
when (val codeGenResult = CodeGenerator.urlCodeGenReader(importDataSchema.url, name, formats, true)) {
174-
is CodeGenerationReadResult.Success -> {
175-
val readDfMethod = codeGenResult.getReadDfMethod(importDataSchema.url.toExternalForm())
176-
val code = readDfMethod.additionalImports.joinToString("\n") +
177-
"\n" +
178-
codeGenResult.code
179-
180-
execute(code)
181-
execute("""DISPLAY("Data schema successfully imported as ${property.name}: $name")""")
182-
183-
name
184-
}
185-
186-
is CodeGenerationReadResult.Error -> {
187-
execute("""DISPLAY("Failed to read data schema from ${importDataSchema.url}: ${codeGenResult.reason}")""")
188-
null
189-
}
245+
updateVariable<Any> { instance, property ->
246+
when (instance) {
247+
is AnyCol -> updateAnyColVariable(instance, property, codeGen)
248+
is ColumnGroup<*> -> updateColumnGroupVariable(instance, property, codeGen)
249+
is AnyRow -> updateAnyRowVariable(instance, property, codeGen)
250+
is AnyFrame -> updateAnyFrameVariable(instance, property, codeGen)
251+
is ImportDataSchema -> updateImportDataSchemaVariable(instance, property)
252+
else -> null
190253
}
191254
}
192255

193-
updateVariable<AnyFrame> { df, property ->
194-
execute(codeGen.process(df, property), property)
195-
}
196-
197-
updateVariable<AnyRow> { row, property ->
198-
execute(codeGen.process(row, property), property)
199-
}
200-
201-
updateVariable<ColumnGroup<*>> { col, property ->
202-
execute(codeGen.process(col.asDataFrame(), property), property)
203-
}
204-
205-
updateVariable<AnyCol> { col, property ->
206-
if (col.isColumnGroup()) {
207-
val codeWithConverter = codeGen.process(col.asColumnGroup().asDataFrame(), property).let { c ->
208-
CodeWithConverter(c.declarations) { c.converter("$it.asColumnGroup()") }
209-
}
210-
execute(codeWithConverter, property)
211-
} else null
212-
}
213-
214256
fun KotlinKernelHost.addDataSchemas(classes: List<KClass<*>>) {
215257
val code = classes.joinToString("\n") {
216258
codeGen.process(it)

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/CodeGenerationTests.kt

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.jetbrains.kotlinx.dataframe.jupyter
22

3+
import org.intellij.lang.annotations.Language
34
import org.jetbrains.kotlinx.jupyter.api.Code
45
import org.junit.Test
56

@@ -11,9 +12,20 @@ class CodeGenerationTests : DataFrameJupyterTest() {
1112
}
1213
}
1314

15+
@Test
16+
fun `Type erased dataframe`() {
17+
@Language("kts")
18+
val a = """
19+
fun create(): Any? = dataFrameOf("a")(1)
20+
val df = create()
21+
df.a
22+
""".checkCompilation()
23+
}
24+
1425
@Test
1526
fun `nullable dataframe`() {
16-
"""
27+
@Language("kts")
28+
val a = """
1729
fun create(): AnyFrame? = dataFrameOf("a")(1)
1830
val df = create()
1931
df.a
@@ -22,7 +34,8 @@ class CodeGenerationTests : DataFrameJupyterTest() {
2234

2335
@Test
2436
fun `nullable columnGroup`() {
25-
"""
37+
@Language("kts")
38+
val a = """
2639
fun create(): AnyCol? = dataFrameOf("a")(1).asColumnGroup().asDataColumn()
2740
val col = create()
2841
col.a
@@ -31,7 +44,8 @@ class CodeGenerationTests : DataFrameJupyterTest() {
3144

3245
@Test
3346
fun `nullable dataRow`() {
34-
"""
47+
@Language("kts")
48+
val a = """
3549
fun create(): AnyRow? = dataFrameOf("a")(1).single()
3650
val row = create()
3751
row.a

‎core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterCodegenTests.kt

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
2626
""".trimIndent()
2727
)
2828
res1 shouldBe Unit
29+
30+
@Language("kts")
2931
val res2 = execRaw("df") as AnyFrame
3032

3133
res2["value"].type shouldBe typeOf<List<Any?>>()
@@ -126,6 +128,7 @@ class JupyterCodegenTests : JupyterReplTestCase() {
126128
)
127129
res1 shouldBe Unit
128130

131+
@Language("kts")
129132
val res2 = execRaw("df.`1`")
130133
res2.shouldBeInstanceOf<ValueColumn<*>>()
131134
}
@@ -141,6 +144,7 @@ class JupyterCodegenTests : JupyterReplTestCase() {
141144
)
142145
res1.shouldBeInstanceOf<MimeTypedResult>()
143146

147+
@Language("kts")
144148
val res2 = exec(
145149
"""listOf(df.`{a}`[0], df.`(b)`[0], df.`{c}`[0])"""
146150
)
@@ -157,6 +161,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
157161
""".trimIndent()
158162
)
159163
res1.shouldBeInstanceOf<MimeTypedResult>()
164+
165+
@Language("kts")
160166
val res2 = exec(
161167
"listOf(df.`\$id`[0])"
162168
)
@@ -174,6 +180,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
174180
)
175181
res1.shouldBeInstanceOf<MimeTypedResult>()
176182
println(res1.entries.joinToString())
183+
184+
@Language("kts")
177185
val res2 = exec(
178186
"listOf(df.`Day's`[0])"
179187
)
@@ -193,6 +201,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
193201
)
194202
res1.shouldBeInstanceOf<MimeTypedResult>()
195203
println(res1.entries.joinToString())
204+
205+
@Language("kts")
196206
val res2 = exec(
197207
"listOf(df.`Test `[0])"
198208
)
@@ -212,6 +222,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
212222
)
213223
res1.shouldBeInstanceOf<MimeTypedResult>()
214224
println(res1.entries.joinToString())
225+
226+
@Language("kts")
215227
val res2 = exec(
216228
"listOf(df.`Test `[0])"
217229
)
@@ -220,6 +232,7 @@ class JupyterCodegenTests : JupyterReplTestCase() {
220232

221233
@Test
222234
fun `generic interface`() {
235+
@Language("kts")
223236
val res1 = exec(
224237
"""
225238
@DataSchema
@@ -229,6 +242,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
229242
""".trimIndent()
230243
)
231244
res1.shouldBeInstanceOf<Unit>()
245+
246+
@Language("kts")
232247
val res2 = exec(
233248
"""
234249
val <T> ColumnsContainer<Generic<T>>.test1: DataColumn<T> get() = field
@@ -240,6 +255,7 @@ class JupyterCodegenTests : JupyterReplTestCase() {
240255

241256
@Test
242257
fun `generic interface with upper bound`() {
258+
@Language("kts")
243259
val res1 = exec(
244260
"""
245261
@DataSchema
@@ -249,6 +265,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
249265
""".trimIndent()
250266
)
251267
res1.shouldBeInstanceOf<Unit>()
268+
269+
@Language("kts")
252270
val res2 = exec(
253271
"""
254272
val <T : String> ColumnsContainer<Generic<T>>.test1: DataColumn<T> get() = field
@@ -260,6 +278,7 @@ class JupyterCodegenTests : JupyterReplTestCase() {
260278

261279
@Test
262280
fun `generic interface with variance and user type in type parameters`() {
281+
@Language("kts")
263282
val res1 = exec(
264283
"""
265284
interface UpperBound
@@ -271,6 +290,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
271290
""".trimIndent()
272291
)
273292
res1.shouldBeInstanceOf<Unit>()
293+
294+
@Language("kts")
274295
val res2 = exec(
275296
"""
276297
val <T : UpperBound> ColumnsContainer<Generic<T>>.test1: DataColumn<T> get() = field
@@ -282,7 +303,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
282303

283304
@Test
284305
fun `generate a new marker when dataframe marker is not a data schema so that columns are accessible with extensions`() {
285-
exec(
306+
@Language("kts")
307+
val a = exec(
286308
"""
287309
enum class State {
288310
Idle, Productive, Maintenance
@@ -305,7 +327,8 @@ class JupyterCodegenTests : JupyterReplTestCase() {
305327
""".trimIndent()
306328
)
307329
shouldNotThrowAny {
308-
exec(
330+
@Language("kts")
331+
val b = exec(
309332
"""
310333
events.toolId
311334
events.state

0 commit comments

Comments
 (0)