feat: add knit_print method for DataFrame (experimental)

DESCRIPTION

@@ -26,7 +26,8 @@ Suggests:
     bit64,
     knitr,
     tibble,
-    rmarkdown
+    rmarkdown,
+    withr
 Config/testthat/edition: 3
 Collate: 
     'utils.R'
@@ -59,6 +60,7 @@ Collate:
     'namespace.R'
     'options.R'
     'parquet.R'
+    'pkg-knitr.R'
     'pkg-nanoarrow.R'
     'rlang.R'
     'rust_result.R'

NAMESPACE

@@ -130,6 +130,7 @@ export(GroupBy_agg)
 export(GroupBy_as_data_frame)
 export(LazyFrame_print)
 export(csv_reader)
+export(knit_print.DataFrame)
 export(ncol.DataFrame)
 export(ncol.LazyFrame)
 export(nrow.DataFrame)

R/pkg-knitr.R

@@ -0,0 +1,23 @@
+#' @export
+knit_print.DataFrame = function(x, ...) {
+  .env_formatting = Sys.getenv("POLARS_FMT_TABLE_FORMATTING")
+  .env_inline_data_type = tolower(Sys.getenv("POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE"))
+
+  on.exit({
+    Sys.setenv(POLARS_FMT_TABLE_FORMATTING = .env_formatting)
+    Sys.setenv(POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE = .env_inline_data_type)
+  })
+
+  if (.env_formatting %in% c("", "ASCII_MARKDOWN") && .env_inline_data_type %in% c("", "true", "1")) {
+    Sys.setenv(POLARS_FMT_TABLE_FORMATTING = "ASCII_MARKDOWN")
+    Sys.setenv(POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE = "1")
+  }
+
+  out <- capture.output(print(x))
+
+  # Needs to insert a blank line
+  out <- c(out[1], "", out[-1]) |>
+    paste(collapse = "\n")
+
+  knitr::asis_output(out)
+}

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")]
 

README.md

@@ -126,24 +126,22 @@ library(polars)
 
 dat = pl$DataFrame(mtcars)
 dat
-#> polars DataFrame: shape: (32, 11)
-#> ┌──────┬─────┬───────┬───────┬─────┬─────┬─────┬──────┬──────┐
-#> │ mpg  ┆ cyl ┆ disp  ┆ hp    ┆ ... ┆ vs  ┆ am  ┆ gear ┆ carb │
-#> │ ---  ┆ --- ┆ ---   ┆ ---   ┆     ┆ --- ┆ --- ┆ ---  ┆ ---  │
-#> │ f64  ┆ f64 ┆ f64   ┆ f64   ┆     ┆ f64 ┆ f64 ┆ f64  ┆ f64  │
-#> ╞══════╪═════╪═══════╪═══════╪═════╪═════╪═════╪══════╪══════╡
-#> │ 21.0 ┆ 6.0 ┆ 160.0 ┆ 110.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 4.0  ┆ 4.0  │
-#> │ 21.0 ┆ 6.0 ┆ 160.0 ┆ 110.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 4.0  ┆ 4.0  │
-#> │ 22.8 ┆ 4.0 ┆ 108.0 ┆ 93.0  ┆ ... ┆ 1.0 ┆ 1.0 ┆ 4.0  ┆ 1.0  │
-#> │ 21.4 ┆ 6.0 ┆ 258.0 ┆ 110.0 ┆ ... ┆ 1.0 ┆ 0.0 ┆ 3.0  ┆ 1.0  │
-#> │ ...  ┆ ... ┆ ...   ┆ ...   ┆ ... ┆ ... ┆ ... ┆ ...  ┆ ...  │
-#> │ 15.8 ┆ 8.0 ┆ 351.0 ┆ 264.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 4.0  │
-#> │ 19.7 ┆ 6.0 ┆ 145.0 ┆ 175.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 6.0  │
-#> │ 15.0 ┆ 8.0 ┆ 301.0 ┆ 335.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 8.0  │
-#> │ 21.4 ┆ 4.0 ┆ 121.0 ┆ 109.0 ┆ ... ┆ 1.0 ┆ 1.0 ┆ 4.0  ┆ 2.0  │
-#> └──────┴─────┴───────┴───────┴─────┴─────┴─────┴──────┴──────┘
 ```
 
+polars DataFrame: shape: (32, 11)
+
+| mpg (f64) | cyl (f64) | disp (f64) | hp (f64) | …   | vs (f64) | am (f64) | gear (f64) | carb (f64) |
+|-----------|-----------|------------|----------|-----|----------|----------|------------|------------|
+| 21.0      | 6.0       | 160.0      | 110.0    | …   | 0.0      | 1.0      | 4.0        | 4.0        |
+| 21.0      | 6.0       | 160.0      | 110.0    | …   | 0.0      | 1.0      | 4.0        | 4.0        |
+| 22.8      | 4.0       | 108.0      | 93.0     | …   | 1.0      | 1.0      | 4.0        | 1.0        |
+| 21.4      | 6.0       | 258.0      | 110.0    | …   | 1.0      | 0.0      | 3.0        | 1.0        |
+| …         | …         | …          | …        | …   | …        | …        | …          | …          |
+| 15.8      | 8.0       | 351.0      | 264.0    | …   | 0.0      | 1.0      | 5.0        | 4.0        |
+| 19.7      | 6.0       | 145.0      | 175.0    | …   | 0.0      | 1.0      | 5.0        | 6.0        |
+| 15.0      | 8.0       | 301.0      | 335.0    | …   | 0.0      | 1.0      | 5.0        | 8.0        |
+| 21.4      | 4.0       | 121.0      | 109.0    | …   | 1.0      | 1.0      | 4.0        | 2.0        |
+
 Once our Polars DataFrame has been created, we can chain together a
 series of data manipulations as part of the same query. For example:
 
@@ -156,19 +154,17 @@ dat$filter(
   pl$col("mpg")$mean()$alias("mean_mpg"),
   pl$col("hp")$median()$alias("med_hp")
 )
-#> polars DataFrame: shape: (4, 4)
-#> ┌─────┬─────┬───────────┬────────┐
-#> │ cyl ┆ am  ┆ mean_mpg  ┆ med_hp │
-#> │ --- ┆ --- ┆ ---       ┆ ---    │
-#> │ f64 ┆ f64 ┆ f64       ┆ f64    │
-#> ╞═════╪═════╪═══════════╪════════╡
-#> │ 6.0 ┆ 1.0 ┆ 20.566667 ┆ 110.0  │
-#> │ 6.0 ┆ 0.0 ┆ 19.125    ┆ 116.5  │
-#> │ 8.0 ┆ 0.0 ┆ 15.05     ┆ 180.0  │
-#> │ 8.0 ┆ 1.0 ┆ 15.4      ┆ 299.5  │
-#> └─────┴─────┴───────────┴────────┘
 ```
 
+polars DataFrame: shape: (4, 4)
+
+| cyl (f64) | am (f64) | mean_mpg (f64) | med_hp (f64) |
+|-----------|----------|----------------|--------------|
+| 6.0       | 1.0      | 20.566667      | 110.0        |
+| 6.0       | 0.0      | 19.125         | 116.5        |
+| 8.0       | 0.0      | 15.05          | 180.0        |
+| 8.0       | 1.0      | 15.4           | 299.5        |
+
 The above is an example of Polars’ eager execution engine. But for
 maximum performance, it is preferable to use Polars’ lazy execution
 mode, which allows the package to apply additional query optimizations.
@@ -184,19 +180,17 @@ ldat$filter(
   pl$col("mpg")$mean()$alias("mean_mpg"),
   pl$col("hp")$median()$alias("med_hp")
 )$collect()
-#> polars DataFrame: shape: (4, 4)
-#> ┌─────┬─────┬───────────┬────────┐
-#> │ cyl ┆ am  ┆ mean_mpg  ┆ med_hp │
-#> │ --- ┆ --- ┆ ---       ┆ ---    │
-#> │ f64 ┆ f64 ┆ f64       ┆ f64    │
-#> ╞═════╪═════╪═══════════╪════════╡
-#> │ 6.0 ┆ 1.0 ┆ 20.566667 ┆ 110.0  │
-#> │ 6.0 ┆ 0.0 ┆ 19.125    ┆ 116.5  │
-#> │ 8.0 ┆ 0.0 ┆ 15.05     ┆ 180.0  │
-#> │ 8.0 ┆ 1.0 ┆ 15.4      ┆ 299.5  │
-#> └─────┴─────┴───────────┴────────┘
 ```
 
+polars DataFrame: shape: (4, 4)
+
+| cyl (f64) | am (f64) | mean_mpg (f64) | med_hp (f64) |
+|-----------|----------|----------------|--------------|
+| 6.0       | 1.0      | 20.566667      | 110.0        |
+| 6.0       | 0.0      | 19.125         | 116.5        |
+| 8.0       | 0.0      | 15.05          | 180.0        |
+| 8.0       | 1.0      | 15.4           | 299.5        |
+
 ## Contribute
 
 Contributions are very welcome!

tests/testthat/_snaps/knitr.md

@@ -0,0 +1,80 @@
+# Snapshot test of knitr
+
+    Code
+      .knit_file("dataframe.Rmd")
+    Output
+
+      ```r
+      pl$DataFrame(mtcars)
+      ```
+
+      polars DataFrame: shape: (32, 11)
+
+      | mpg (f64) | cyl (f64) | disp (f64) | hp (f64) | ... | vs (f64) | am (f64) | gear (f64) | carb (f64) |
+      |-----------|-----------|------------|----------|-----|----------|----------|------------|------------|
+      | 21.0      | 6.0       | 160.0      | 110.0    | ... | 0.0      | 1.0      | 4.0        | 4.0        |
+      | 21.0      | 6.0       | 160.0      | 110.0    | ... | 0.0      | 1.0      | 4.0        | 4.0        |
+      | 22.8      | 4.0       | 108.0      | 93.0     | ... | 1.0      | 1.0      | 4.0        | 1.0        |
+      | 21.4      | 6.0       | 258.0      | 110.0    | ... | 1.0      | 0.0      | 3.0        | 1.0        |
+      | ...       | ...       | ...        | ...      | ... | ...      | ...      | ...        | ...        |
+      | 15.8      | 8.0       | 351.0      | 264.0    | ... | 0.0      | 1.0      | 5.0        | 4.0        |
+      | 19.7      | 6.0       | 145.0      | 175.0    | ... | 0.0      | 1.0      | 5.0        | 6.0        |
+      | 15.0      | 8.0       | 301.0      | 335.0    | ... | 0.0      | 1.0      | 5.0        | 8.0        |
+      | 21.4      | 4.0       | 121.0      | 109.0    | ... | 1.0      | 1.0      | 4.0        | 2.0        |
+
+---
+
+    Code
+      .knit_file("dataframe.Rmd")
+    Output
+
+      ```r
+      pl$DataFrame(mtcars)
+      ```
+
+      polars DataFrame: shape: (32, 11)
+
+      ┌──────┬─────┬───────┬───────┬─────┬─────┬─────┬──────┬──────┐
+      │ mpg  ┆ cyl ┆ disp  ┆ hp    ┆ ... ┆ vs  ┆ am  ┆ gear ┆ carb │
+      │ ---  ┆ --- ┆ ---   ┆ ---   ┆     ┆ --- ┆ --- ┆ ---  ┆ ---  │
+      │ f64  ┆ f64 ┆ f64   ┆ f64   ┆     ┆ f64 ┆ f64 ┆ f64  ┆ f64  │
+      ╞══════╪═════╪═══════╪═══════╪═════╪═════╪═════╪══════╪══════╡
+      │ 21.0 ┆ 6.0 ┆ 160.0 ┆ 110.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 4.0  ┆ 4.0  │
+      │ 21.0 ┆ 6.0 ┆ 160.0 ┆ 110.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 4.0  ┆ 4.0  │
+      │ 22.8 ┆ 4.0 ┆ 108.0 ┆ 93.0  ┆ ... ┆ 1.0 ┆ 1.0 ┆ 4.0  ┆ 1.0  │
+      │ 21.4 ┆ 6.0 ┆ 258.0 ┆ 110.0 ┆ ... ┆ 1.0 ┆ 0.0 ┆ 3.0  ┆ 1.0  │
+      │ ...  ┆ ... ┆ ...   ┆ ...   ┆ ... ┆ ... ┆ ... ┆ ...  ┆ ...  │
+      │ 15.8 ┆ 8.0 ┆ 351.0 ┆ 264.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 4.0  │
+      │ 19.7 ┆ 6.0 ┆ 145.0 ┆ 175.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 6.0  │
+      │ 15.0 ┆ 8.0 ┆ 301.0 ┆ 335.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 8.0  │
+      │ 21.4 ┆ 4.0 ┆ 121.0 ┆ 109.0 ┆ ... ┆ 1.0 ┆ 1.0 ┆ 4.0  ┆ 2.0  │
+      └──────┴─────┴───────┴───────┴─────┴─────┴─────┴──────┴──────┘
+
+---
+
+    Code
+      .knit_file("dataframe.Rmd")
+    Output
+
+      ```r
+      pl$DataFrame(mtcars)
+      ```
+
+      polars DataFrame: shape: (32, 11)
+
+      ┌──────┬─────┬───────┬───────┬─────┬─────┬─────┬──────┬──────┐
+      │ mpg  ┆ cyl ┆ disp  ┆ hp    ┆ ... ┆ vs  ┆ am  ┆ gear ┆ carb │
+      │ ---  ┆ --- ┆ ---   ┆ ---   ┆     ┆ --- ┆ --- ┆ ---  ┆ ---  │
+      │ f64  ┆ f64 ┆ f64   ┆ f64   ┆     ┆ f64 ┆ f64 ┆ f64  ┆ f64  │
+      ╞══════╪═════╪═══════╪═══════╪═════╪═════╪═════╪══════╪══════╡
+      │ 21.0 ┆ 6.0 ┆ 160.0 ┆ 110.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 4.0  ┆ 4.0  │
+      │ 21.0 ┆ 6.0 ┆ 160.0 ┆ 110.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 4.0  ┆ 4.0  │
+      │ 22.8 ┆ 4.0 ┆ 108.0 ┆ 93.0  ┆ ... ┆ 1.0 ┆ 1.0 ┆ 4.0  ┆ 1.0  │
+      │ 21.4 ┆ 6.0 ┆ 258.0 ┆ 110.0 ┆ ... ┆ 1.0 ┆ 0.0 ┆ 3.0  ┆ 1.0  │
+      │ ...  ┆ ... ┆ ...   ┆ ...   ┆ ... ┆ ... ┆ ... ┆ ...  ┆ ...  │
+      │ 15.8 ┆ 8.0 ┆ 351.0 ┆ 264.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 4.0  │
+      │ 19.7 ┆ 6.0 ┆ 145.0 ┆ 175.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 6.0  │
+      │ 15.0 ┆ 8.0 ┆ 301.0 ┆ 335.0 ┆ ... ┆ 0.0 ┆ 1.0 ┆ 5.0  ┆ 8.0  │
+      │ 21.4 ┆ 4.0 ┆ 121.0 ┆ 109.0 ┆ ... ┆ 1.0 ┆ 1.0 ┆ 4.0  ┆ 2.0  │
+      └──────┴─────┴───────┴───────┴─────┴─────┴─────┴──────┴──────┘
+

tests/testthat/files/dataframe.Rmd

@@ -0,0 +1,3 @@
+```{r}
+pl$DataFrame(mtcars)
+```

tests/testthat/test-knitr.R

@@ -0,0 +1,23 @@
+.knit_file <- function(file_name) {
+  file <- file.path("files", file_name)
+  output <- tempfile(fileext = "md")
+  on.exit(unlink(output))
+
+  suppressWarnings(knitr::knit(file, output, quiet = TRUE, envir = new.env()))
+
+  readLines(output) |>
+    paste0(collapse = "\n") |>
+    cat()
+}
+
+test_that("Snapshot test of knitr", {
+  expect_snapshot(.knit_file("dataframe.Rmd"), cran = TRUE)
+  withr::with_envvar(
+    new = c("POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE" = "false"),
+    expect_snapshot(.knit_file("dataframe.Rmd"), cran = TRUE)
+  )
+  withr::with_envvar(
+    new = c("POLARS_FMT_TABLE_FORMATTING" = "DEFAULT"),
+    expect_snapshot(.knit_file("dataframe.Rmd"), cran = TRUE)
+  )
+})