feat: add knit_print method for DataFrame (experimental)

DESCRIPTION

@@ -27,6 +27,7 @@ Suggests:
     bit64,
     knitr,
     tibble,
+    pillar,
     rmarkdown,
     withr
 Config/Needs/website:
@@ -69,6 +70,7 @@ Collate:
     'namespace.R'
     'options.R'
     'parquet.R'
+    'pkg-knitr.R'
     'pkg-nanoarrow.R'
     'rlang.R'
     'rust_result.R'

NAMESPACE

@@ -135,6 +135,7 @@ export(GroupBy_agg)
 export(GroupBy_to_data_frame)
 export(LazyFrame_print)
 export(csv_reader)
+export(knit_print.DataFrame)
 export(ncol.DataFrame)
 export(nrow.DataFrame)
 export(pl)
@@ -146,6 +147,7 @@ importFrom(stats,na.omit)
 importFrom(utils,.DollarNames)
 importFrom(utils,capture.output)
 importFrom(utils,download.file)
+importFrom(utils,getFromNamespace)
 importFrom(utils,globalVariables)
 importFrom(utils,head)
 importFrom(utils,str)

NEWS.md

@@ -2,6 +2,8 @@
 
 ## What's changed
 - `DataFrame` objects can be subsetted using brackets like standard R data frames: `pl$DataFrame(mtcars)[2:4, c("mpg", "hp")]` (#140 @vincentarelbundock)
+- An experimental `knit_print()` method has been added to DataFrame that outputs HTML tables
+  (similar to py-polars' HTML output) (#125 @eitsupi)
 - `Series` gains new methods: `$mean`, `$median`, `$std`, `$var` (#170 @vincentarelbundock)
 
 # polars v0.5.0

R/pkg-knitr.R

@@ -0,0 +1,177 @@
+#' knit print polars DataFrame
+#'
+#' Mimics python-polars' NotebookFormatter
+#' for HTML outputs.
+#'
+#' Outputs HTML tables if the output format is HTML
+#' and the document's `df_print` option is not `"default"` or `"tibble"`.
+#'
+#' Or, the output format can be enforced with R's `options` function as follows:
+#'
+#' - `options(polars.df_print = "default")` for the default print method.
+#' - `options(polars.df_print = "html")` for the HTML table.
+#' @name knit_print.DataFrame
+#' @param x a polars DataFrame to knit_print
+#' @param ... additional arguments, not used
+#' @keywords DataFrame
+#' @export
+knit_print.DataFrame = function(x, ...) {
+  .print_opt = getOption("polars.df_print", "auto")
+  .rmd_df_print = knitr::opts_knit$get("rmarkdown.df_print")
+  if (isTRUE(.print_opt == "html") ||
+    (isTRUE(.print_opt != "default") &&
+      !isTRUE(.rmd_df_print %in% c("default", "tibble")) &&
+      knitr::is_html_output())) {
+    x |>
+      to_html_table() |>
+      knitr::raw_html() |>
+      knitr::asis_output()
+  } else {
+    print(x)
+  }
+}
+
+#' Generate HTML table from a DataFrame
+#' @param x DataFrame
+#' @param max_cols an integer of maximum number of columns to display
+#' @param max_rows an integer of maximum number of rows to display
+#' @return A string of HTML code
+#' @examples
+#' to_html_table(mtcars, 3, 3)
+#' @noRd
+#' @importFrom utils getFromNamespace
+to_html_table = function(x, max_cols = 75, max_rows = 40) {
+  if (!requireNamespace("knitr", quietly = TRUE)) {
+    stop("Please install the `knitr` package to use `to_html_table`.")
+  }
+
+  escape_html = getFromNamespace("escape_html", "knitr")
+  omit_chr = "…"
+
+  .dim = dim(x)
+
+  df_height = .dim[1]
+  df_width = .dim[2]
+
+  row_idx = .idx(max_rows, df_height)
+  col_idx = .idx(max_cols, df_width)
+
+  cols = names(x)[col_idx]
+  col_names = cols |>
+    escape_html()
+  col_types = x |>
+    .get_dtype_strings() |>
+    escape_html()
+
+  if (max_cols <= df_width) {
+    col_names = .cut_off(col_names, max_cols, omit_chr)
+    col_types = .cut_off(col_types, max_cols, omit_chr)
+  }
+
+  .header_names = col_names |>
+    .tag("th") |>
+    .tag("tr")
+
+  .header_dtypes = col_types |>
+    .tag("td") |>
+    .tag("tr")
+
+  .header_all = .header_names |>
+    paste0(.header_dtypes) |>
+    .tag("thead")
+
+  .env_str_len = Sys.getenv("POLARS_FMT_STR_LEN")
+  .str_len = ifelse(.env_str_len == "", 15, as.integer(.env_str_len))
+
+  chr_mat = sapply(cols, \(col) as.character(x[row_idx, col, drop = TRUE], str_length = .str_len)) |>
+    escape_html() |>
+    matrix(nrow = length(row_idx))
+
+  if (max_cols <= df_width) {
+    .seq = seq_along(cols)
+    chr_mat = cbind(
+      chr_mat[, head(.seq, max_cols %/% 2)],
+      omit_chr,
+      chr_mat[, tail(.seq, max_cols %/% 2)]
+    )
+  }
+
+  if (max_rows <= df_height) {
+    .seq = seq_along(row_idx)
+    chr_mat = rbind(
+      chr_mat[head(.seq, max_rows %/% 2), ],
+      omit_chr,
+      chr_mat[tail(.seq, max_rows %/% 2), ]
+    )
+  }
+
+  .body = chr_mat |>
+    t() |>
+    as.data.frame() |>
+    sapply(\(x) .tag(x, "td")) |>
+    .tag("tr") |>
+    .tag("tbody")
+
+  .style = "<style>
+.dataframe > thead > tr > th,
+.dataframe > tbody > tr > td {
+  text-align: right;
+}
+</style>
+"
+
+  .shape = sprintf("<small>shape: (%s, %s)</small>", .dim[1], .dim[2])
+
+  paste0(.header_all, .body) |>
+    .tag("table", c(border = "1", class = "dataframe")) |>
+    (\(x) (paste0(.style, .shape, x)))() |>
+    .tag("div")
+}
+
+.idx = function(.max, .length) {
+  if (.max <= .length) {
+    out = c(seq_len(.max %/% 2), seq(.length - .max %/% 2 + 1L, .length))
+  } else {
+    out = seq_len(.length)
+  }
+  out
+}
+
+.cut_off = function(.vec, .max, omit_chr) {
+  c(head(.vec, .max %/% 2), omit_chr, tail(.vec, .max %/% 2))
+}
+
+#' @param .elements chr vector
+#' @param .tag single chr
+#' @param .attr named chr vector
+#' @return single charactor
+#' @examples
+#' .tag(letters[1:2], "tr")
+#' @noRd
+.tag = function(.elements, .tag, .attr = NULL) {
+  if (!is.null(.attr)) {
+    .pre = paste0("<", .tag, " ", paste0(sprintf('%s="%s"', names(.attr), .attr), collapse = " "), ">")
+  } else {
+    .pre = c(paste0("<", .tag, ">"))
+  }
+
+  .post = paste0("</", .tag, ">")
+
+  paste0(.pre, .elements, .post, collapse = "")
+}
+
+#' Get types of each column
+#' @param df DataFrame like object
+#' @return string vector of column type names
+#' @noRd
+.get_dtype_strings = function(df) {
+  if (inherits(df, "DataFrame")) {
+    df$dtype_strings()
+  } else {
+    if (!requireNamespace("pillar", quietly = TRUE)) {
+      stop("Please install the `pillar` package to use `to_html_table` for non-polars objects.")
+    }
+    sapply(names(df), \(x) pillar::type_sum(df[, x, drop = TRUE])) |>
+      unname()
+  }
+}

R/s3_methods.R

@@ -171,20 +171,19 @@ max.LazyFrame = function(x, ...) x$max()
 #' @noRd
 as.vector.Series = function(x, mode) x$to_vector()
 
-#' as.character for Series
+#' as.character for polars Series
 #' @param x Series
-#' @param format a logical. If `TRUE`, the Series will be formatted.
-#' @param str_length an integer. If `format = TRUE`,
+#' @param str_length an integer. If specified,
+#' @param ... Additional arguments are ignored.
 #' utf8 or categorical type Series will be formatted to a string of this length.
 #' @param ... Additional characters are ignored.
 #' @examples
 #' s = pl$Series(c("foo", "barbaz"))
 #' as.character(s)
-#' as.character(s, format = TRUE)
-#' as.character(s, format = TRUE, str_length = 3)
+#' as.character(s, str_length = 3)
 #' @export
-as.character.Series = function(x, ..., format = FALSE, str_length = 15) {
-    if (isTRUE(format)) {
+as.character.Series = function(x, ..., str_length = NULL) {
+    if (is.numeric(str_length) && str_length > 0) {
         .pr$Series$to_fmt_char(x, str_length = str_length)
     } else {
         x$to_vector() |>

R/zzz.R

@@ -137,6 +137,7 @@ pl$mem_address = mem_address
   s3_register("nanoarrow::infer_nanoarrow_schema", "DataFrame")
   s3_register("arrow::as_record_batch_reader", "DataFrame")
   s3_register("arrow::as_arrow_table", "DataFrame")
+  s3_register("knitr::knit_print", "DataFrame")
 
   pl$numeric_dtypes = pl$dtypes[substr(names(pl$dtypes),1,3) %in% c("Int","Flo")]