Merge pull request #1459 from Kotlin/formatHeader

Format header

Author: AndreiKingsley

core/api/core.api

@@ -2232,11 +2232,21 @@ public final class org/jetbrains/kotlinx/dataframe/api/ForEachKt {
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/FormatClause {
-	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;)V
-	public synthetic fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;)V
+	public synthetic fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class org/jetbrains/kotlinx/dataframe/api/FormatHeaderKt {
+	public static final fun formatHeader (Lorg/jetbrains/kotlinx/dataframe/DataFrame;)Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;
+	public static final fun formatHeader (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;
+	public static final fun formatHeader (Lorg/jetbrains/kotlinx/dataframe/DataFrame;[Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;
+	public static final fun formatHeader (Lorg/jetbrains/kotlinx/dataframe/api/FormattedFrame;)Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;
+	public static final fun formatHeader (Lorg/jetbrains/kotlinx/dataframe/api/FormattedFrame;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;
+	public static final fun formatHeader (Lorg/jetbrains/kotlinx/dataframe/api/FormattedFrame;[Ljava/lang/String;)Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;
+	public static final fun with (Lorg/jetbrains/kotlinx/dataframe/api/HeaderFormatClause;Lkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/api/FormattedFrame;
+}
+
 public final class org/jetbrains/kotlinx/dataframe/api/FormatKt {
 	public static final fun and (Lorg/jetbrains/kotlinx/dataframe/api/CellAttributes;Lorg/jetbrains/kotlinx/dataframe/api/CellAttributes;)Lorg/jetbrains/kotlinx/dataframe/api/CellAttributes;
 	public static final fun at (Lorg/jetbrains/kotlinx/dataframe/api/FormatClause;Ljava/util/Collection;)Lorg/jetbrains/kotlinx/dataframe/api/FormatClause;
@@ -2259,8 +2269,8 @@ public final class org/jetbrains/kotlinx/dataframe/api/FormatKt {
 }
 
 public final class org/jetbrains/kotlinx/dataframe/api/FormattedFrame {
-	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function3;)V
-	public synthetic fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function3;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;)V
+	public synthetic fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun getDisplayConfiguration (Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;)Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;
 	public final fun toHtml (Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;)Lorg/jetbrains/kotlinx/dataframe/io/DataFrameHtmlData;
 	public static synthetic fun toHtml$default (Lorg/jetbrains/kotlinx/dataframe/api/FormattedFrame;Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;ILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/io/DataFrameHtmlData;
@@ -2545,6 +2555,12 @@ public final class org/jetbrains/kotlinx/dataframe/api/HeadKt {
 	public static synthetic fun head$default (Lorg/jetbrains/kotlinx/dataframe/DataFrame;IILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
 }
 
+public final class org/jetbrains/kotlinx/dataframe/api/HeaderFormatClause {
+	public fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;)V
+	public synthetic fun <init> (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function3;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun toString ()Ljava/lang/String;
+}
+
 public final class org/jetbrains/kotlinx/dataframe/api/ImplodeKt {
 	public static final fun implode (Lorg/jetbrains/kotlinx/dataframe/DataFrame;Z)Lorg/jetbrains/kotlinx/dataframe/DataRow;
 	public static final fun implode (Lorg/jetbrains/kotlinx/dataframe/DataFrame;ZLkotlin/jvm/functions/Function2;)Lorg/jetbrains/kotlinx/dataframe/DataFrame;
@@ -6205,23 +6221,29 @@ public final class org/jetbrains/kotlinx/dataframe/io/DisplayConfiguration {
 	public static final field Companion Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration$Companion;
 	public synthetic fun <init> (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Ljava/lang/String;ZZZZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public synthetic fun <init> (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Ljava/lang/String;ZZZZZLkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public synthetic fun <init> (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;Ljava/lang/String;ZZZZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public synthetic fun <init> (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;Ljava/lang/String;ZZZZZLkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun component1 ()Ljava/lang/Integer;
 	public final fun component10 ()Z
+	public final fun component11 ()Z
 	public final fun component2 ()Ljava/lang/Integer;
 	public final fun component3 ()I
 	public final fun component4 ()Lkotlin/jvm/functions/Function3;
-	public final fun component5-3Sl7FsM ()Ljava/lang/String;
-	public final fun component6 ()Z
-	public final fun component8 ()Z
+	public final fun component5 ()Lkotlin/jvm/functions/Function2;
+	public final fun component6-3Sl7FsM ()Ljava/lang/String;
+	public final fun component7 ()Z
 	public final fun component9 ()Z
-	public final fun copy-rqXL5tM (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Ljava/lang/String;ZZZZZ)Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;
+	public final fun copy-bMNacXk (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;Ljava/lang/String;ZZZZZ)Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;
+	public static synthetic fun copy-bMNacXk$default (Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Lkotlin/jvm/functions/Function2;Ljava/lang/String;ZZZZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;
+	public final synthetic fun copy-rqXL5tM (Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Ljava/lang/String;ZZZZZ)Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;
 	public static synthetic fun copy-rqXL5tM$default (Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;Ljava/lang/Integer;Ljava/lang/Integer;ILkotlin/jvm/functions/Function3;Ljava/lang/String;ZZZZZILjava/lang/Object;)Lorg/jetbrains/kotlinx/dataframe/io/DisplayConfiguration;
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getCellContentLimit ()I
 	public final fun getCellFormatter ()Lkotlin/jvm/functions/Function3;
 	public final fun getDecimalFormat-3Sl7FsM ()Ljava/lang/String;
 	public final fun getDownsizeBufferedImage ()Z
 	public final fun getEnableFallbackStaticTables ()Z
+	public final fun getHeaderFormatter ()Lkotlin/jvm/functions/Function2;
 	public final fun getIsolatedOutputs ()Z
 	public final fun getNestedRowsLimit ()Ljava/lang/Integer;
 	public final fun getRowsLimit ()Ljava/lang/Integer;
@@ -6233,6 +6255,7 @@ public final class org/jetbrains/kotlinx/dataframe/io/DisplayConfiguration {
 	public final fun setDecimalFormat-h5o3lmc (Ljava/lang/String;)V
 	public final fun setDownsizeBufferedImage (Z)V
 	public final fun setEnableFallbackStaticTables (Z)V
+	public final fun setHeaderFormatter (Lkotlin/jvm/functions/Function2;)V
 	public final fun setIsolatedOutputs (Z)V
 	public final fun setNestedRowsLimit (Ljava/lang/Integer;)V
 	public final fun setRowsLimit (Ljava/lang/Integer;)V

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

@@ -325,7 +325,7 @@ public fun <T, C> DataFrame<T>.format(vararg columns: KProperty<C>): FormatClaus
  *   If unspecified, all columns will be formatted.
  */
 public fun <T, C> FormattedFrame<T>.format(columns: ColumnsSelector<T, C>): FormatClause<T, C> =
-    FormatClause(df, columns, formatter)
+    FormatClause(df, columns, formatter, oldHeaderFormatter = headerFormatter)
 
 /**
  * @include [CommonFormatDocs]
@@ -390,7 +390,13 @@ public fun <T> FormattedFrame<T>.format(): FormatClause<T, Any?> = FormatClause(
  * Check out the full [Grammar][FormatDocs.Grammar].
  */
 public fun <T, C> FormatClause<T, C>.where(filter: RowValueFilter<T, C>): FormatClause<T, C> =
-    FormatClause(filter = this.filter and filter, df = df, columns = columns, oldFormatter = oldFormatter)
+    FormatClause(
+        filter = this.filter and filter,
+        df = df,
+        columns = columns,
+        oldFormatter = oldFormatter,
+        oldHeaderFormatter = oldHeaderFormatter,
+    )
 
 /**
  * Only format the selected columns at given row indices.
@@ -780,7 +786,11 @@ public typealias CellFormatter<C> = FormattingDsl.(cell: C) -> CellAttributes?
  *
  * You can apply further formatting to this [FormattedFrame] by calling [format()][FormattedFrame.format] once again.
  */
-public class FormattedFrame<T>(internal val df: DataFrame<T>, internal val formatter: RowColFormatter<T, *>? = null) {
+public class FormattedFrame<T>(
+    internal val df: DataFrame<T>,
+    internal val formatter: RowColFormatter<T, *>? = null,
+    internal val headerFormatter: HeaderColFormatter<*>? = null,
+) {
 
     /**
      * Returns a [DataFrameHtmlData] without additional definitions.
@@ -826,7 +836,10 @@ public class FormattedFrame<T>(internal val df: DataFrame<T>, internal val forma
     /** Applies this formatter to the given [configuration] and returns a new instance. */
     @Suppress("UNCHECKED_CAST")
     public fun getDisplayConfiguration(configuration: DisplayConfiguration): DisplayConfiguration =
-        configuration.copy(cellFormatter = formatter as RowColFormatter<*, *>?)
+        configuration.copy(
+            cellFormatter = formatter as RowColFormatter<*, *>?,
+            headerFormatter = headerFormatter,
+        )
 }
 
 /**
@@ -858,6 +871,7 @@ public class FormatClause<T, C>(
     internal val columns: ColumnsSelector<T, C> = { all().cast() },
     internal val oldFormatter: RowColFormatter<T, C>? = null,
     internal val filter: RowValueFilter<T, C> = { true },
+    internal val oldHeaderFormatter: HeaderColFormatter<*>? = null,
 ) {
     override fun toString(): String =
         "FormatClause(df=$df, columns=$columns, oldFormatter=$oldFormatter, filter=$filter)"

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

@@ -0,0 +1,143 @@
+package org.jetbrains.kotlinx.dataframe.api
+
+import org.jetbrains.kotlinx.dataframe.ColumnsSelector
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
+import org.jetbrains.kotlinx.dataframe.impl.api.formatHeaderImpl
+
+// region docs
+
+/**
+ * A lambda used to format a column header (its displayed name) when rendering a dataframe to HTML.
+ *
+ * The lambda runs in the context of [FormattingDsl] and receives the [ColumnWithPath] of the header to format.
+ * Return a [CellAttributes] (or `null`) describing the CSS you want to apply to the header cell.
+ *
+ * Examples:
+ * - Center the header: `attr("text-align", "center")`
+ * - Make it bold: `bold`
+ * - Set a custom color: `textColor(rgb(10, 10, 10))`
+ */
+public typealias HeaderColFormatter<C> = FormattingDsl.(col: ColumnWithPath<C>) -> CellAttributes?
+
+/**
+ * An intermediate class used in the header-format operation [formatHeader].
+ *
+ * This class itself does nothing—it represents a selection of columns whose headers will be formatted.
+ * Finalize this step by calling [with] to produce a new [FormattedFrame].
+ *
+ * Header formatting is additive and supports nested column groups: styles specified for a parent group
+ * are inherited by its child columns unless overridden for the child.
+ */
+public class HeaderFormatClause<T, C>(
+    internal val df: DataFrame<T>,
+    internal val columns: ColumnsSelector<T, C> = { all().cast() },
+    internal val oldHeaderFormatter: HeaderColFormatter<C>? = null,
+    internal val oldCellFormatter: RowColFormatter<T, *>? = null,
+) {
+    override fun toString(): String =
+        "HeaderFormatClause(df=$df, columns=$columns, oldHeaderFormatter=$oldHeaderFormatter, oldCellFormatter=$oldCellFormatter)"
+}
+
+// endregion
+
+// region DataFrame.formatHeader
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Selects [columns] whose headers should be formatted. If unspecified, all columns will be formatted.
+ *
+ * This does not immediately produce a [FormattedFrame]; instead it returns a [HeaderFormatClause]
+ * which must be finalized using [HeaderFormatClause.with].
+ */
+public fun <T, C> DataFrame<T>.formatHeader(columns: ColumnsSelector<T, C>): HeaderFormatClause<T, C> =
+    HeaderFormatClause(this, columns)
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Selects [columns] whose headers should be formatted.
+ *
+ * This does not immediately produce a [FormattedFrame]; instead it returns a [HeaderFormatClause]
+ * which must be finalized using [HeaderFormatClause.with].
+ */
+public fun <T> DataFrame<T>.formatHeader(vararg columns: String): HeaderFormatClause<T, Any?> =
+    formatHeader { columns.toColumnSet() }
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Selects all columns for header formatting.
+ *
+ * This does not immediately produce a [FormattedFrame]; instead it returns a [HeaderFormatClause]
+ * which must be finalized using [HeaderFormatClause.with].
+ */
+public fun <T> DataFrame<T>.formatHeader(): HeaderFormatClause<T, Any?> = HeaderFormatClause(this)
+
+// endregion
+
+// region FormattedFrame.formatHeader
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Selects [columns] whose headers should be formatted.
+ *
+ * This does not immediately produce a [FormattedFrame]; instead it returns a [HeaderFormatClause]
+ * which must be finalized using [HeaderFormatClause.with].
+ */
+public fun <T, C> FormattedFrame<T>.formatHeader(columns: ColumnsSelector<T, C>): HeaderFormatClause<T, C> =
+    HeaderFormatClause(
+        df = df,
+        columns = columns,
+        oldHeaderFormatter = headerFormatter as HeaderColFormatter<C>?,
+        oldCellFormatter = formatter,
+    )
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Selects [columns] whose headers should be formatted.
+ *
+ * This does not immediately produce a [FormattedFrame]; instead it returns a [HeaderFormatClause]
+ * which must be finalized using [HeaderFormatClause.with].
+ */
+public fun <T> FormattedFrame<T>.formatHeader(vararg columns: String): HeaderFormatClause<T, Any?> =
+    formatHeader { columns.toColumnSet() }
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Selects all columns for header formatting.
+ *
+ * This does not immediately produce a [FormattedFrame]; instead it returns a [HeaderFormatClause]
+ * which must be finalized using [HeaderFormatClause.with].
+ */
+public fun <T> FormattedFrame<T>.formatHeader(): HeaderFormatClause<T, Any?> =
+    HeaderFormatClause(
+        df = df,
+        oldHeaderFormatter = headerFormatter,
+        oldCellFormatter = formatter,
+    )
+
+// endregion
+
+// region terminal operations
+
+/**
+ * **Experimental API. It may be changed in the future.**
+ *
+ * Creates a new [FormattedFrame] that uses the specified [HeaderColFormatter] to format the selected headers.
+ *
+ * Header formatting is additive: attributes from already-applied header formatters are combined with the newly
+ *
+ *  returned attributes using [CellAttributes.and]. If a parent column group is selected, its attributes are
+ * applied to its children unless explicitly overridden.
+ */
+@Suppress("UNCHECKED_CAST")
+public fun <T, C> HeaderFormatClause<T, C>.with(formatter: HeaderColFormatter<C>): FormattedFrame<T> =
+    formatHeaderImpl(formatter)
+
+// endregion

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

@@ -56,15 +56,18 @@ internal inline fun <T, C> FormatClause<T, C>.formatImpl(
     val clause = this
     val columns = clause.df.getColumnPaths(UnresolvedColumnsPolicy.Skip, clause.columns).toSet()
 
-    return FormattedFrame(clause.df) { row, col ->
-        val oldAttributes = clause.oldFormatter?.invoke(FormattingDsl, row, col.cast())
-        if (col.path in columns) {
-            val value = col[row] as C
-            if (clause.filter(row, value)) {
-                return@FormattedFrame oldAttributes and formatter(FormattingDsl, row.cast(), col.cast())
+    return FormattedFrame(
+        df = clause.df,
+        formatter = { row, col ->
+            val oldAttributes = clause.oldFormatter?.invoke(FormattingDsl, row, col.cast())
+            if (col.path in columns) {
+                val value = col[row] as C
+                if (clause.filter(row, value)) {
+                    return@FormattedFrame oldAttributes and formatter(FormattingDsl, row.cast(), col.cast())
+                }
             }
-        }
-
-        oldAttributes
-    }
+            oldAttributes
+        },
+        headerFormatter = clause.oldHeaderFormatter,
+    )
 }

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

@@ -0,0 +1,24 @@
+package org.jetbrains.kotlinx.dataframe.impl.api
+
+import org.jetbrains.kotlinx.dataframe.api.FormattedFrame
+import org.jetbrains.kotlinx.dataframe.api.FormattingDsl
+import org.jetbrains.kotlinx.dataframe.api.HeaderColFormatter
+import org.jetbrains.kotlinx.dataframe.api.HeaderFormatClause
+import org.jetbrains.kotlinx.dataframe.api.and
+import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
+import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
+import org.jetbrains.kotlinx.dataframe.impl.getColumnPaths
+
+internal fun <T, C> HeaderFormatClause<T, C>.formatHeaderImpl(formatter: HeaderColFormatter<C>): FormattedFrame<T> {
+    val selectedPaths = df.getColumnPaths(UnresolvedColumnsPolicy.Skip, columns).toSet()
+    val oldHeader = oldHeaderFormatter
+
+    val composedHeader: HeaderColFormatter<Any?> = { col ->
+        val typedCol = col as ColumnWithPath<C>
+        val existingAttr = oldHeader?.invoke(FormattingDsl, typedCol)
+        val newAttr = if (col.path in selectedPaths) formatter(FormattingDsl, typedCol) else null
+        existingAttr and newAttr
+    }
+
+    return FormattedFrame(df, oldCellFormatter, composedHeader)
+}