Fix bug in code generation for nullable variables.

Author: nikitinas

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/Integration.kt

@@ -119,7 +119,10 @@ internal class Integration : JupyterIntegration() {
             } else null
         }
 
-        fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, property: KProperty<*>): VariableName? = execute(codeWithConverter, property.name)
+        fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, property: KProperty<*>): VariableName? {
+            val variableName = property.name + if (property.returnType.isMarkedNullable) "!!" else ""
+            return execute(codeWithConverter, variableName)
+        }
 
         updateVariable<AnyFrame> { df, property ->
             execute(codeGen.process(df, property), property)
@@ -135,8 +138,10 @@ internal class Integration : JupyterIntegration() {
 
         updateVariable<AnyCol> { col, property ->
             if (col.isColumnGroup()) {
-                val codeWithConverter = codeGen.process(col.asColumnGroup().asDataFrame(), property)
-                execute(codeWithConverter, "${property.name}.asColumnGroup()")
+                val codeWithConverter = codeGen.process(col.asColumnGroup().asDataFrame(), property).let { c ->
+                    CodeWithConverter(c.declarations) { c.converter(it + ".asColumnGroup()") }
+                }
+                execute(codeWithConverter, property)
             } else null
         }
 

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/CodeGenerationTests.kt

@@ -0,0 +1,40 @@
+package org.jetbrains.kotlinx.dataframe.jupyter
+
+import org.jetbrains.kotlinx.jupyter.api.Code
+import org.junit.Test
+
+class CodeGenerationTests : DataFrameJupyterTest() {
+
+    private fun Code.checkCompilation() {
+        lines().forEach {
+            exec(it)
+        }
+    }
+
+    @Test
+    fun `nullable dataframe`() {
+        """
+            fun create(): AnyFrame? = dataFrameOf("a")(1)
+            val df = create()
+            df.a
+        """.checkCompilation()
+    }
+
+    @Test
+    fun `nullable columnGroup`() {
+        """
+            fun create(): AnyCol? = dataFrameOf("a")(1).asColumnGroup().asDataColumn()
+            val col = create()
+            col.a
+        """.checkCompilation()
+    }
+
+    @Test
+    fun `nullable dataRow`() {
+        """
+            fun create(): AnyRow? = dataFrameOf("a")(1).single()
+            val row = create()
+            row.a
+        """.checkCompilation()
+    }
+}