GroupBy Jupyter codegen

core/api/core.api

@@ -673,6 +673,7 @@ public final class org/jetbrains/kotlinx/dataframe/api/CastKt {
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/DataColumn;)Lorg/jetbrains/kotlinx/dataframe/DataColumn;
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/DataFrame;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/DataRow;)Lorg/jetbrains/kotlinx/dataframe/DataRow;
+	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/api/GroupBy;)Lorg/jetbrains/kotlinx/dataframe/api/GroupBy;
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/columns/ColumnAccessor;)Lorg/jetbrains/kotlinx/dataframe/columns/ColumnAccessor;
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/columns/ColumnGroup;)Lorg/jetbrains/kotlinx/dataframe/columns/ColumnGroup;
 	public static final fun cast (Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;)Lorg/jetbrains/kotlinx/dataframe/columns/ColumnReference;
@@ -4770,23 +4771,23 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/codeGen/BaseFiel
 }
 
 public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult {
-	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;Ljava/util/List;)V
-	public final fun component1 ()Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;Ljava/util/List;)V
+	public final fun component1 ()Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 	public final fun component2 ()Ljava/util/List;
-	public final fun copy (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;Ljava/util/List;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;
-	public static synthetic fun copy$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;Ljava/util/List;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;
+	public final fun copy (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;Ljava/util/List;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;
+	public static synthetic fun copy$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;Ljava/util/List;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;
 	public fun equals (Ljava/lang/Object;)Z
-	public final fun getCode ()Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public final fun getCode ()Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 	public final fun getNewMarkers ()Ljava/util/List;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
 }
 
 public abstract interface class org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator : org/jetbrains/kotlinx/dataframe/codeGen/ExtensionsCodeGenerator {
 	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator$Companion;
-	public abstract fun generate (Lorg/jetbrains/kotlinx/dataframe/codeGen/Marker;Lorg/jetbrains/kotlinx/dataframe/codeGen/InterfaceGenerationMode;ZLorg/jetbrains/kotlinx/dataframe/codeGen/DefaultReadDfMethod;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public abstract fun generate (Lorg/jetbrains/kotlinx/dataframe/codeGen/Marker;Lorg/jetbrains/kotlinx/dataframe/codeGen/InterfaceGenerationMode;ZLorg/jetbrains/kotlinx/dataframe/codeGen/DefaultReadDfMethod;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 	public abstract fun generate (Lorg/jetbrains/kotlinx/dataframe/schema/DataFrameSchema;Ljava/lang/String;ZZZLorg/jetbrains/kotlinx/dataframe/codeGen/MarkerVisibility;Ljava/lang/Iterable;Lorg/jetbrains/kotlinx/dataframe/codeGen/DefaultReadDfMethod;Lorg/jetbrains/kotlinx/dataframe/codeGen/NameNormalizer;Z)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;
-	public static synthetic fun generate$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator;Lorg/jetbrains/kotlinx/dataframe/codeGen/Marker;Lorg/jetbrains/kotlinx/dataframe/codeGen/InterfaceGenerationMode;ZLorg/jetbrains/kotlinx/dataframe/codeGen/DefaultReadDfMethod;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public static synthetic fun generate$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator;Lorg/jetbrains/kotlinx/dataframe/codeGen/Marker;Lorg/jetbrains/kotlinx/dataframe/codeGen/InterfaceGenerationMode;ZLorg/jetbrains/kotlinx/dataframe/codeGen/DefaultReadDfMethod;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 	public static synthetic fun generate$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator;Lorg/jetbrains/kotlinx/dataframe/schema/DataFrameSchema;Ljava/lang/String;ZZZLorg/jetbrains/kotlinx/dataframe/codeGen/MarkerVisibility;Ljava/lang/Iterable;Lorg/jetbrains/kotlinx/dataframe/codeGen/DefaultReadDfMethod;Lorg/jetbrains/kotlinx/dataframe/codeGen/NameNormalizer;ZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;
 }
 
@@ -4796,30 +4797,30 @@ public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator$Compani
 }
 
 public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeGeneratorKt {
-	public static final fun generate (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator;Lkotlin/reflect/KClass;Lorg/jetbrains/kotlinx/dataframe/codeGen/InterfaceGenerationMode;Z)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public static final fun generate (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator;Lkotlin/reflect/KClass;Lorg/jetbrains/kotlinx/dataframe/codeGen/InterfaceGenerationMode;Z)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 }
 
-public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter {
-	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter$Companion;
+public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator {
+	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator$Companion;
 	public static final field EMPTY_DECLARATIONS Ljava/lang/String;
-	public fun <init> (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
-	public synthetic fun <init> (Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Ljava/lang/String;Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;)V
+	public synthetic fun <init> (Ljava/lang/String;Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()Ljava/lang/String;
-	public final fun component2 ()Lkotlin/jvm/functions/Function1;
-	public final fun copy (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
-	public static synthetic fun copy$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public final fun component2 ()Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;
+	public final fun copy (Ljava/lang/String;Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
+	public static synthetic fun copy$default (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;Ljava/lang/String;Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
+	public final fun declarationsWithCastExpression (Ljava/lang/String;)Ljava/lang/String;
 	public fun equals (Ljava/lang/Object;)Z
-	public final fun getConverter ()Lkotlin/jvm/functions/Function1;
 	public final fun getDeclarations ()Ljava/lang/String;
-	public final fun getHasConverter ()Z
+	public final fun getHasCaster ()Z
 	public final fun getHasDeclarations ()Z
+	public final fun getTypeCastGenerator ()Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;
 	public fun hashCode ()I
 	public fun toString ()Ljava/lang/String;
-	public final fun with (Ljava/lang/String;)Ljava/lang/String;
 }
 
-public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter$Companion {
-	public final fun getEMPTY ()Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+public final class org/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator$Companion {
+	public final fun getEMPTY ()Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 	public final fun getEMPTY_CONVERTER ()Lkotlin/jvm/functions/Function1;
 }
 
@@ -4830,7 +4831,7 @@ public abstract interface class org/jetbrains/kotlinx/dataframe/codeGen/DefaultR
 
 public abstract interface class org/jetbrains/kotlinx/dataframe/codeGen/ExtensionsCodeGenerator {
 	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/codeGen/ExtensionsCodeGenerator$Companion;
-	public abstract fun generate (Lorg/jetbrains/kotlinx/dataframe/codeGen/IsolatedMarker;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;
+	public abstract fun generate (Lorg/jetbrains/kotlinx/dataframe/codeGen/IsolatedMarker;)Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;
 }
 
 public final class org/jetbrains/kotlinx/dataframe/codeGen/ExtensionsCodeGenerator$Companion {
@@ -4976,6 +4977,28 @@ public final class org/jetbrains/kotlinx/dataframe/codeGen/NameNormalizer : kotl
 public final class org/jetbrains/kotlinx/dataframe/codeGen/NameNormalizer$Companion {
 }
 
+public abstract interface class org/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator {
+	public abstract fun addCastTo (Ljava/lang/String;)Ljava/lang/String;
+	public fun invoke (Ljava/lang/String;)Ljava/lang/String;
+}
+
+public final class org/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator$DataFrameApi : org/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator {
+	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator$DataFrameApi$Companion;
+	public synthetic fun <init> ([Ljava/lang/String;Lkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun addCastTo (Ljava/lang/String;)Ljava/lang/String;
+	public final fun getTypes ()[Ljava/lang/String;
+}
+
+public final class org/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator$DataFrameApi$Companion {
+	public final fun invoke ([Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator;
+}
+
+public final class org/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator$Empty : org/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator {
+	public static final field INSTANCE Lorg/jetbrains/kotlinx/dataframe/codeGen/TypeCastGenerator$Empty;
+	public fun addCastTo (Ljava/lang/String;)Ljava/lang/String;
+	public fun invoke (Ljava/lang/String;)Ljava/lang/String;
+}
+
 public final class org/jetbrains/kotlinx/dataframe/codeGen/ValidFieldName {
 	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/codeGen/ValidFieldName$Companion;
 	public synthetic fun <init> (Ljava/lang/String;ZLkotlin/jvm/internal/DefaultConstructorMarker;)V
@@ -5850,7 +5873,7 @@ public final class org/jetbrains/kotlinx/dataframe/impl/codeGen/CodeGeneratorImp
 	public static final fun quoteIfNeeded (Ljava/lang/String;)Ljava/lang/String;
 	public static final fun toStandaloneSnippet (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;)Ljava/lang/String;
 	public static final fun toStandaloneSnippet (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeGenResult;Ljava/lang/String;Ljava/util/List;)Ljava/lang/String;
-	public static final fun toStandaloneSnippet (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithConverter;Ljava/lang/String;Ljava/util/List;)Ljava/lang/String;
+	public static final fun toStandaloneSnippet (Lorg/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator;Ljava/lang/String;Ljava/util/List;)Ljava/lang/String;
 }
 
 public abstract interface class org/jetbrains/kotlinx/dataframe/impl/codeGen/DfReadResult {

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/cast.kt

@@ -99,3 +99,5 @@ public fun <C> TransformableColumnSet<*>.cast(): TransformableColumnSet<C> = thi
 public fun <C> TransformableSingleColumn<*>.cast(): TransformableSingleColumn<C> = this as TransformableSingleColumn<C>
 
 public fun <C> ColumnReference<*>.cast(): ColumnReference<C> = this as ColumnReference<C>
+
+public fun <T, G> GroupBy<*, *>.cast(): GroupBy<T, G> = this as GroupBy<T, G>

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator.kt

@@ -15,7 +15,7 @@ public enum class InterfaceGenerationMode {
     None,
 }
 
-public data class CodeGenResult(val code: CodeWithConverter, val newMarkers: List<Marker>)
+public data class CodeGenResult(val code: CodeWithTypeCastGenerator, val newMarkers: List<Marker>)
 
 public interface CodeGenerator : ExtensionsCodeGenerator {
 
@@ -37,7 +37,7 @@ public interface CodeGenerator : ExtensionsCodeGenerator {
         interfaceMode: InterfaceGenerationMode,
         extensionProperties: Boolean,
         readDfMethod: DefaultReadDfMethod? = null,
-    ): CodeWithConverter
+    ): CodeWithTypeCastGenerator
 
     public companion object {
         public fun create(useFqNames: Boolean = true): CodeGenerator =
@@ -54,7 +54,7 @@ internal fun CodeGenerator.generate(
     markerClass: KClass<*>,
     interfaceMode: InterfaceGenerationMode,
     extensionProperties: Boolean,
-): CodeWithConverter =
+): CodeWithTypeCastGenerator =
     generate(
         MarkersExtractor.get(markerClass),
         interfaceMode,
@@ -64,4 +64,4 @@ internal fun CodeGenerator.generate(
 public inline fun <reified T> CodeGenerator.generate(
     interfaceMode: InterfaceGenerationMode,
     extensionProperties: Boolean,
-): CodeWithConverter = generate(T::class, interfaceMode, extensionProperties)
+): CodeWithTypeCastGenerator = generate(T::class, interfaceMode, extensionProperties)

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeWithTypeCastGenerator.kt

@@ -0,0 +1,81 @@
+package org.jetbrains.kotlinx.dataframe.codeGen
+
+import org.jetbrains.kotlinx.dataframe.api.cast
+
+/**
+ * Class representing generated code declarations for a [Marker].
+ *
+ * @param declarations The generated code.
+ * @param typeCastGenerator Optional [TypeCastGenerator] for the [Marker], see [TypeCastGenerator] for
+ *   more information.
+ */
+public data class CodeWithTypeCastGenerator(
+    val declarations: Code,
+    val typeCastGenerator: TypeCastGenerator = TypeCastGenerator.Empty,
+) {
+
+    public companion object {
+        public const val EMPTY_DECLARATIONS: Code = ""
+
+        @Deprecated("", ReplaceWith("TypeCastGenerator.Empty"))
+        public val EMPTY_CONVERTER: (Expression) -> Code = { it }
+        public val EMPTY: CodeWithTypeCastGenerator = CodeWithTypeCastGenerator(EMPTY_DECLARATIONS)
+    }
+
+    val hasDeclarations: Boolean get() = declarations.isNotBlank()
+
+    val hasCaster: Boolean
+        get() = typeCastGenerator !is TypeCastGenerator.Empty &&
+            typeCastGenerator("it").trim() != "it"
+
+    public fun declarationsWithCastExpression(expression: Expression): Code =
+        when {
+            !hasCaster -> declarations
+            !hasDeclarations -> typeCastGenerator(expression)
+            else -> declarations + "\n" + typeCastGenerator(expression)
+        }
+}
+
+public typealias Code = String
+public typealias Expression = String
+
+/**
+ * A [TypeCastGenerator] can generate [Code] given an [Expression] that casts or converts
+ * it to a predefined target type.
+ *
+ * To create a [TypeCastGenerator] that, for instance, casts everything you pass to [Any?][Any]:
+ * ```kt
+ * val myCastGenerator = TypeCastGenerator { expression -> "(\$expression as Any?)" }
+ *
+ * myCastGenerator.addCastTo("myVariable") == "(myVariable as Any?)"
+ * ```
+ *
+ * @see TypeCastGenerator.Empty
+ * @see TypeCastGenerator.DataFrameApi
+ */
+public fun interface TypeCastGenerator {
+    public fun addCastTo(expression: Expression): Code
+
+    public operator fun invoke(expression: Expression): Code = addCastTo(expression)
+
+    public object Empty : TypeCastGenerator by (TypeCastGenerator { it })
+
+    /**
+     * [TypeCastGenerator] that uses the [cast] functions of the DataFrame API.
+     *
+     * NOTE: This generator assumes there's a `.cast<>()` function available that can be called on your
+     * specific [Expression]. It will cause runtime errors when there isn't one.
+     */
+    public class DataFrameApi private constructor(public val types: Array<out String>) : TypeCastGenerator {
+        override fun addCastTo(expression: Expression): Code =
+            if (types.isEmpty()) {
+                "$expression.cast()"
+            } else {
+                "$expression.cast<${types.joinToString()}>()"
+            }
+
+        public companion object {
+            public operator fun invoke(vararg types: String): TypeCastGenerator = DataFrameApi(types)
+        }
+    }
+}

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/ExtensionsCodeGenerator.kt

@@ -4,7 +4,7 @@ import org.jetbrains.kotlinx.dataframe.impl.codeGen.ExtensionsCodeGeneratorImpl
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.ShortNames
 
 public interface ExtensionsCodeGenerator {
-    public fun generate(marker: IsolatedMarker): CodeWithConverter
+    public fun generate(marker: IsolatedMarker): CodeWithTypeCastGenerator
 
     public companion object {
         public fun create(): ExtensionsCodeGenerator = ExtensionsCodeGeneratorImpl(ShortNames)