Skip to content

Commit fe6716e

Browse files
committed
stable "shuffle" test
this is a verbose variant of manual generation, should be improved at some point
1 parent 7f3ba71 commit fe6716e

File tree

4 files changed

+125
-122
lines changed

4 files changed

+125
-122
lines changed

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/explainer/PluginCallback.kt

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

3-
import com.beust.klaxon.JsonObject
43
import org.jetbrains.kotlinx.dataframe.AnyCol
54
import org.jetbrains.kotlinx.dataframe.AnyFrame
65
import org.jetbrains.kotlinx.dataframe.DataFrame
@@ -22,7 +21,6 @@ import org.jetbrains.kotlinx.dataframe.api.Update
2221
import org.jetbrains.kotlinx.dataframe.api.format
2322
import org.jetbrains.kotlinx.dataframe.api.frames
2423
import org.jetbrains.kotlinx.dataframe.api.into
25-
import org.jetbrains.kotlinx.dataframe.api.print
2624
import org.jetbrains.kotlinx.dataframe.api.toDataFrame
2725
import org.jetbrains.kotlinx.dataframe.api.values
2826
import org.jetbrains.kotlinx.dataframe.api.where
@@ -35,71 +33,6 @@ import org.jetbrains.kotlinx.dataframe.io.tableInSessionId
3533
import org.jetbrains.kotlinx.dataframe.io.toHTML
3634
import java.io.File
3735

38-
private fun convertToHTML(dataframeLike: Any): DataFrameHtmlData {
39-
fun DataFrame<*>.toHTML() = toHTML(DisplayConfiguration(), getFooter = { "" })
40-
fun FormattedFrame<*>.toHTML1() = toHTML(DisplayConfiguration())
41-
42-
return when (dataframeLike) {
43-
is Pivot<*> -> dataframeLike.frames().toDataFrame().toHTML()
44-
is ReducedPivot<*> -> dataframeLike.values().toDataFrame().toHTML()
45-
is PivotGroupBy<*> -> dataframeLike.frames().toHTML()
46-
is ReducedPivotGroupBy<*> -> dataframeLike.values().toHTML()
47-
is SplitWithTransform<*, *, *> -> dataframeLike.into().toHTML()
48-
is Merge<*, *, *> -> dataframeLike.into("merged").toHTML()
49-
is Gather<*, *, *, *> -> dataframeLike.into("key", "value").toHTML()
50-
is Update<*, *> -> dataframeLike.df.let {
51-
var it = it.format(dataframeLike.columns as ColumnsSelectionDsl<Any?>.(it: ColumnsSelectionDsl<Any?>) -> ColumnSet<*>)
52-
if (dataframeLike.filter != null) {
53-
it = it.where(dataframeLike.filter as RowValueFilter<Any?, Any?>)
54-
}
55-
it.with {
56-
background(rgb(152, 251, 152))
57-
}
58-
}
59-
.toHTML1()
60-
is Convert<*, *> -> DataFrameHtmlData(body = "<p>${dataframeLike::class}</p>")
61-
is FormattedFrame<*> -> dataframeLike.toHTML1()
62-
is GroupBy<*, *> -> dataframeLike.toDataFrame().toHTML()
63-
is AnyFrame -> dataframeLike.toHTML()
64-
is AnyCol -> dataframeLike.toDataFrame().toHTML()
65-
is DataRow<*> -> dataframeLike.toDataFrame().toHTML()
66-
is Split<*, *> -> dataframeLike.toDataFrame().toHTML()
67-
// is MoveClause<*, *>-> null
68-
// is RenameClause<*, *> -> null
69-
// is ReplaceClause<*, *> -> null
70-
// is GroupClause<*, *> -> null
71-
// is InsertClause<*> -> null
72-
// is FormatClause<*, *> -> null
73-
else -> throw IllegalArgumentException("Unsupported type: ${dataframeLike::class}")
74-
}
75-
}
76-
77-
private fun convertToDescription(dataframeLike: Any): String {
78-
return when (dataframeLike) {
79-
is AnyFrame -> dataframeLike.let { "DataFrame: rowsCount = ${it.rowsCount()}, columnsCount = ${it.columnsCount()}" }
80-
is Pivot<*> -> "Pivot"
81-
is ReducedPivot<*> -> "ReducedPivot"
82-
is PivotGroupBy<*> -> "PivotGroupBy"
83-
is ReducedPivotGroupBy<*> -> "ReducedPivotGroupBy"
84-
is SplitWithTransform<*, *, *> -> "SplitWithTransform"
85-
is Split<*, *> -> "Split"
86-
// is MoveClause<*, *> -> "Move"
87-
// is RenameClause<*, *> -> "Rename"
88-
// is ReplaceClause<*, *> -> "Replace"
89-
// is GroupClause<*, *> -> "Group"
90-
// is InsertClause<*> -> "Insert"
91-
// is FormatClause<*, *> -> "Format"
92-
is Merge<*, *, *> -> "Merge"
93-
is Gather<*, *, *, *> -> "Gather"
94-
is Update<*, *> -> "Update"
95-
is Convert<*, *> -> "Convert"
96-
is FormattedFrame<*> -> "FormattedFrame"
97-
is GroupBy<*, *> -> "GroupBy"
98-
is DataRow<*> -> "DataRow"
99-
else -> "TODO"
100-
}.escapeHtmlForIFrame()
101-
}
102-
10336
annotation class TransformDataFrameExpressions
10437

10538
object PluginCallback {
@@ -122,6 +55,7 @@ object PluginCallback {
12255
sessionId = 0
12356
tableInSessionId = 0
12457
var output = DataFrameHtmlData.tableDefinitions() + DataFrameHtmlData(
58+
// copy writerside stlyles
12559
style = """
12660
body {
12761
font-family: "JetBrains Mono",SFMono-Regular,Consolas,"Liberation Mono",Menlo,Courier,monospace;
@@ -150,13 +84,13 @@ object PluginCallback {
15084
// make copy to avoid concurrent modification exception
15185
val statements = expressionsByStatement.toMap()
15286
when (statements.size) {
153-
0 -> TODO("wtf")
87+
0 -> TODO("function doesn't have any dataframe expression")
15488
1 -> {
155-
output += expressionOutputs(statements.values.single(), open = false)
89+
output += statementOutput(statements.values.single())
15690
}
15791
else -> {
15892
statements.forEach { (index, expressions) ->
159-
var details: DataFrameHtmlData = expressionOutputs(expressions, open = true)
93+
var details: DataFrameHtmlData = statementOutput(expressions)
16094

16195
details = details.copy(
16296
body =
@@ -185,12 +119,9 @@ object PluginCallback {
185119
output.writeHTML(File(destination, "$name.html"))
186120
}
187121

188-
private fun expressionOutputs(
122+
private fun statementOutput(
189123
expressions: List<Expression>,
190-
open: Boolean,
191124
): DataFrameHtmlData {
192-
// val attribute = if (open) " open" else ""
193-
val attribute = ""
194125
var data = DataFrameHtmlData()
195126
if (expressions.size < 2) error("Sample without output or input (i.e. function returns some value)")
196127
for ((i, expression) in expressions.withIndex()) {
@@ -199,7 +130,7 @@ object PluginCallback {
199130
val table = convertToHTML(expression.df)
200131
val description = table.copy(
201132
body = """
202-
<details$attribute>
133+
<details>
203134
<summary>Input ${convertToDescription(expression.df)}</summary>
204135
${table.body}
205136
</details>
@@ -212,7 +143,7 @@ object PluginCallback {
212143
val table = convertToHTML(expression.df)
213144
val description = table.copy(
214145
body = """
215-
<details$attribute>
146+
<details>
216147
<summary>Output ${convertToDescription(expression.df)}</summary>
217148
${table.body}
218149
</details>
@@ -244,41 +175,6 @@ object PluginCallback {
244175
val element = Expression(source, containingClassFqName, containingFunName, df)
245176
list?.plus(element) ?: listOf(element)
246177
}
247-
// strings.add(string)
248-
// names.add(name)
249-
// Can be called with the same name multiple times, need to aggregate samples by function name somehow?
250-
// save schema
251-
val path = "$containingClassFqName.$containingFunName.html"
252-
// names.compute(path) { }
253-
// dfs.add(path)
254-
if (df is AnyFrame) {
255-
println(source)
256-
// df.print()
257-
println(id)
258-
println(receiverId)
259-
} else {
260-
println(df::class)
261-
}
262-
File("build/out").let {
263-
val json = JsonObject(
264-
mapOf(
265-
"string" to source,
266-
"name" to name,
267-
"path" to path,
268-
"id" to id,
269-
"receiverId" to receiverId,
270-
)
271-
).toJsonString()
272-
it.appendText(json)
273-
it.appendText(",\n")
274-
}
275-
println(path)
276-
if (df is AnyFrame) {
277-
df.print()
278-
} else {
279-
println(df::class)
280-
}
281-
// convertToHTML(df).writeHTML(File("build/dataframes/$path"))
282178
}
283179

284180
@Suppress("unused")
@@ -296,6 +192,71 @@ object PluginCallback {
296192
}
297193
}
298194

195+
private fun convertToHTML(dataframeLike: Any): DataFrameHtmlData {
196+
fun DataFrame<*>.toHTML() = toHTML(DisplayConfiguration(), getFooter = { "" })
197+
fun FormattedFrame<*>.toHTML1() = toHTML(DisplayConfiguration())
198+
199+
return when (dataframeLike) {
200+
is Pivot<*> -> dataframeLike.frames().toDataFrame().toHTML()
201+
is ReducedPivot<*> -> dataframeLike.values().toDataFrame().toHTML()
202+
is PivotGroupBy<*> -> dataframeLike.frames().toHTML()
203+
is ReducedPivotGroupBy<*> -> dataframeLike.values().toHTML()
204+
is SplitWithTransform<*, *, *> -> dataframeLike.into().toHTML()
205+
is Merge<*, *, *> -> dataframeLike.into("merged").toHTML()
206+
is Gather<*, *, *, *> -> dataframeLike.into("key", "value").toHTML()
207+
is Update<*, *> -> dataframeLike.df.let {
208+
var it = it.format(dataframeLike.columns as ColumnsSelectionDsl<Any?>.(it: ColumnsSelectionDsl<Any?>) -> ColumnSet<*>)
209+
if (dataframeLike.filter != null) {
210+
it = it.where(dataframeLike.filter as RowValueFilter<Any?, Any?>)
211+
}
212+
it.with {
213+
background(rgb(152, 251, 152))
214+
}
215+
}
216+
.toHTML1()
217+
is Convert<*, *> -> DataFrameHtmlData(body = "<p>${dataframeLike::class}</p>")
218+
is FormattedFrame<*> -> dataframeLike.toHTML1()
219+
is GroupBy<*, *> -> dataframeLike.toDataFrame().toHTML()
220+
is AnyFrame -> dataframeLike.toHTML()
221+
is AnyCol -> dataframeLike.toDataFrame().toHTML()
222+
is DataRow<*> -> dataframeLike.toDataFrame().toHTML()
223+
is Split<*, *> -> dataframeLike.toDataFrame().toHTML()
224+
// is MoveClause<*, *>-> null
225+
// is RenameClause<*, *> -> null
226+
// is ReplaceClause<*, *> -> null
227+
// is GroupClause<*, *> -> null
228+
// is InsertClause<*> -> null
229+
// is FormatClause<*, *> -> null
230+
else -> throw IllegalArgumentException("Unsupported type: ${dataframeLike::class}")
231+
}
232+
}
233+
234+
private fun convertToDescription(dataframeLike: Any): String {
235+
return when (dataframeLike) {
236+
is AnyFrame -> dataframeLike.let { "DataFrame: rowsCount = ${it.rowsCount()}, columnsCount = ${it.columnsCount()}" }
237+
is Pivot<*> -> "Pivot"
238+
is ReducedPivot<*> -> "ReducedPivot"
239+
is PivotGroupBy<*> -> "PivotGroupBy"
240+
is ReducedPivotGroupBy<*> -> "ReducedPivotGroupBy"
241+
is SplitWithTransform<*, *, *> -> "SplitWithTransform"
242+
is Split<*, *> -> "Split"
243+
// is MoveClause<*, *> -> "Move"
244+
// is RenameClause<*, *> -> "Rename"
245+
// is ReplaceClause<*, *> -> "Replace"
246+
// is GroupClause<*, *> -> "Group"
247+
// is InsertClause<*> -> "Insert"
248+
// is FormatClause<*, *> -> "Format"
249+
is Merge<*, *, *> -> "Merge"
250+
is Gather<*, *, *, *> -> "Gather"
251+
is Update<*, *> -> "Update"
252+
is Convert<*, *> -> "Convert"
253+
is FormattedFrame<*> -> "FormattedFrame"
254+
is GroupBy<*, *> -> "GroupBy"
255+
is DataRow<*> -> "DataRow"
256+
else -> throw IllegalArgumentException("Unsupported type: ${dataframeLike::class}")
257+
}.escapeHtmlForIFrame()
258+
}
259+
299260
internal fun String.escapeHtmlForIFrame(): String {
300261
val str = this
301262
return buildString {

core/generated-sources/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Modify.kt

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import org.jetbrains.kotlinx.dataframe.api.fillNulls
3333
import org.jetbrains.kotlinx.dataframe.api.filter
3434
import org.jetbrains.kotlinx.dataframe.api.flatten
3535
import org.jetbrains.kotlinx.dataframe.api.gather
36+
import org.jetbrains.kotlinx.dataframe.api.getRows
3637
import org.jetbrains.kotlinx.dataframe.api.group
3738
import org.jetbrains.kotlinx.dataframe.api.groupBy
3839
import org.jetbrains.kotlinx.dataframe.api.gt
@@ -97,8 +98,10 @@ import org.jetbrains.kotlinx.dataframe.api.with
9798
import org.jetbrains.kotlinx.dataframe.api.withNull
9899
import org.jetbrains.kotlinx.dataframe.api.withValue
99100
import org.jetbrains.kotlinx.dataframe.api.withZero
101+
import org.jetbrains.kotlinx.dataframe.explainer.PluginCallback
100102
import org.jetbrains.kotlinx.dataframe.explainer.TransformDataFrameExpressions
101103
import org.jetbrains.kotlinx.dataframe.impl.api.mapNotNullValues
104+
import org.jetbrains.kotlinx.dataframe.indices
102105
import org.jetbrains.kotlinx.dataframe.io.readJsonStr
103106
import org.jetbrains.kotlinx.dataframe.io.renderToString
104107
import org.jetbrains.kotlinx.dataframe.testResource
@@ -251,11 +254,25 @@ class Modify : TestBase() {
251254
}
252255

253256
@Test
254-
@TransformDataFrameExpressions
255257
fun shuffle() {
256258
// SampleStart
257259
df.shuffle()
258260
// SampleEnd
261+
262+
PluginCallback.expressionsByStatement[0] = listOf(
263+
PluginCallback.Expression(
264+
source = "df",
265+
containingClassFqName = "org.jetbrains.kotlinx.dataframe.samples.api.Modify",
266+
containingFunName = "shuffle",
267+
df = df
268+
),
269+
PluginCallback.Expression(
270+
source = "shuffle()",
271+
containingClassFqName = "org.jetbrains.kotlinx.dataframe.samples.api.Modify",
272+
containingFunName = "shuffle",
273+
df = df.getRows(df.indices.shuffled(Random(123)))
274+
)
275+
)
259276
}
260277

261278
@Test
@@ -898,13 +915,17 @@ class Modify : TestBase() {
898915
}
899916

900917
private class CityInfo(val city: String?, val population: Int, val location: String)
901-
private fun queryCityInfo(city: String?): CityInfo { return CityInfo(city, city?.length ?: 0, "35.5 32.2") }
918+
919+
private fun queryCityInfo(city: String?): CityInfo {
920+
return CityInfo(city, city?.length ?: 0, "35.5 32.2")
921+
}
902922

903923
@Test
904924
@TransformDataFrameExpressions
905925
fun addCalculatedApi() {
906926
// SampleStart
907927
class CityInfo(val city: String?, val population: Int, val location: String)
928+
908929
fun queryCityInfo(city: String?): CityInfo {
909930
return CityInfo(city, city?.length ?: 0, "35.5 32.2")
910931
}

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/Modify.kt

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import org.jetbrains.kotlinx.dataframe.api.fillNulls
3333
import org.jetbrains.kotlinx.dataframe.api.filter
3434
import org.jetbrains.kotlinx.dataframe.api.flatten
3535
import org.jetbrains.kotlinx.dataframe.api.gather
36+
import org.jetbrains.kotlinx.dataframe.api.getRows
3637
import org.jetbrains.kotlinx.dataframe.api.group
3738
import org.jetbrains.kotlinx.dataframe.api.groupBy
3839
import org.jetbrains.kotlinx.dataframe.api.gt
@@ -97,8 +98,10 @@ import org.jetbrains.kotlinx.dataframe.api.with
9798
import org.jetbrains.kotlinx.dataframe.api.withNull
9899
import org.jetbrains.kotlinx.dataframe.api.withValue
99100
import org.jetbrains.kotlinx.dataframe.api.withZero
101+
import org.jetbrains.kotlinx.dataframe.explainer.PluginCallback
100102
import org.jetbrains.kotlinx.dataframe.explainer.TransformDataFrameExpressions
101103
import org.jetbrains.kotlinx.dataframe.impl.api.mapNotNullValues
104+
import org.jetbrains.kotlinx.dataframe.indices
102105
import org.jetbrains.kotlinx.dataframe.io.readJsonStr
103106
import org.jetbrains.kotlinx.dataframe.io.renderToString
104107
import org.jetbrains.kotlinx.dataframe.testResource
@@ -251,11 +254,25 @@ class Modify : TestBase() {
251254
}
252255

253256
@Test
254-
@TransformDataFrameExpressions
255257
fun shuffle() {
256258
// SampleStart
257259
df.shuffle()
258260
// SampleEnd
261+
262+
PluginCallback.expressionsByStatement[0] = listOf(
263+
PluginCallback.Expression(
264+
source = "df",
265+
containingClassFqName = "org.jetbrains.kotlinx.dataframe.samples.api.Modify",
266+
containingFunName = "shuffle",
267+
df = df
268+
),
269+
PluginCallback.Expression(
270+
source = "shuffle()",
271+
containingClassFqName = "org.jetbrains.kotlinx.dataframe.samples.api.Modify",
272+
containingFunName = "shuffle",
273+
df = df.getRows(df.indices.shuffled(Random(123)))
274+
)
275+
)
259276
}
260277

261278
@Test
@@ -898,13 +915,17 @@ class Modify : TestBase() {
898915
}
899916

900917
private class CityInfo(val city: String?, val population: Int, val location: String)
901-
private fun queryCityInfo(city: String?): CityInfo { return CityInfo(city, city?.length ?: 0, "35.5 32.2") }
918+
919+
private fun queryCityInfo(city: String?): CityInfo {
920+
return CityInfo(city, city?.length ?: 0, "35.5 32.2")
921+
}
902922

903923
@Test
904924
@TransformDataFrameExpressions
905925
fun addCalculatedApi() {
906926
// SampleStart
907927
class CityInfo(val city: String?, val population: Int, val location: String)
928+
908929
fun queryCityInfo(city: String?): CityInfo {
909930
return CityInfo(city, city?.length ?: 0, "35.5 32.2")
910931
}

0 commit comments

Comments
 (0)