add S3 methods for tidymodels objects

DESCRIPTION

@@ -23,6 +23,7 @@ Imports:
     cli,
     dplyr,
     generics,
+    hardhat,
     partykit,
     purrr,
     rlang (>= 1.1.0),
@@ -47,6 +48,7 @@ Suggests:
     spelling,
     testthat (>= 3.0.0),
     tidyr,
+    workflows,
     xgboost
 Config/roxygen2/version: 8.1.0
 Config/Needs/website: tidyverse/tidytemplate, ranger, palmerpenguins, tidyr

NAMESPACE

@@ -9,19 +9,23 @@ S3method(active_predictors,cforest)
 S3method(active_predictors,cubist)
 S3method(active_predictors,grf)
 S3method(active_predictors,lgb.Booster)
+S3method(active_predictors,model_fit)
 S3method(active_predictors,party)
 S3method(active_predictors,randomForest)
 S3method(active_predictors,ranger)
 S3method(active_predictors,rpart)
+S3method(active_predictors,workflow)
 S3method(active_predictors,xgb.Booster)
 S3method(as.party,C5.0)
 S3method(as.party,bagger)
 S3method(as.party,bart)
 S3method(as.party,catboost.Model)
 S3method(as.party,grf)
 S3method(as.party,lgb.Booster)
+S3method(as.party,model_fit)
 S3method(as.party,randomForest)
 S3method(as.party,ranger)
+S3method(as.party,workflow)
 S3method(as.party,xgb.Booster)
 S3method(extract_rules,C5.0)
 S3method(extract_rules,ObliqueForest)
@@ -32,10 +36,12 @@ S3method(extract_rules,cforest)
 S3method(extract_rules,cubist)
 S3method(extract_rules,grf)
 S3method(extract_rules,lgb.Booster)
+S3method(extract_rules,model_fit)
 S3method(extract_rules,party)
 S3method(extract_rules,randomForest)
 S3method(extract_rules,ranger)
 S3method(extract_rules,rpart)
+S3method(extract_rules,workflow)
 S3method(extract_rules,xgb.Booster)
 S3method(num_ensemble_members,C5.0)
 S3method(num_ensemble_members,ObliqueForest)
@@ -46,18 +52,22 @@ S3method(num_ensemble_members,cforest)
 S3method(num_ensemble_members,cubist)
 S3method(num_ensemble_members,grf)
 S3method(num_ensemble_members,lgb.Booster)
+S3method(num_ensemble_members,model_fit)
 S3method(num_ensemble_members,randomForest)
 S3method(num_ensemble_members,ranger)
+S3method(num_ensemble_members,workflow)
 S3method(num_ensemble_members,xgb.Booster)
 S3method(var_imp,ObliqueForest)
 S3method(var_imp,catboost.Model)
 S3method(var_imp,cforest)
 S3method(var_imp,grf)
 S3method(var_imp,lgb.Booster)
+S3method(var_imp,model_fit)
 S3method(var_imp,party)
 S3method(var_imp,randomForest)
 S3method(var_imp,ranger)
 S3method(var_imp,rpart)
+S3method(var_imp,workflow)
 S3method(var_imp,xgb.Booster)
 export(active_predictors)
 export(as.party)

NEWS.md

@@ -1,5 +1,7 @@
 # lorax (development version)
 
+* `active_predictors()`, `as.party()`, `extract_rules()`, `num_ensemble_members()`, and `var_imp()` now have methods for fitted parsnip `model_fit` and workflows `workflow` objects. These extract the engine model with `extract_fit_engine()` and pass other arguments to the engine's method.
+
 * `num_ensemble_members()` is a new generic that returns the number of trees (or boosting trials or committees) in an ensemble model.
 
 * `active_predictors()`, `as.party()`, `extract_rules()`, and `var_imp()` now have methods for catboost models (`catboost.Model`). Symmetric and non-symmetric trees are supported. Categorical splits, both one-hot and those based on target statistics (CTRs), are converted to sets of factor levels, including CTRs on combinations of features.

R/lorax-package.R

@@ -14,7 +14,7 @@ generics::var_imp
 
 ## usethis namespace: end
 
-utils::globalVariables("id")
+utils::globalVariables(c("id", "rules"))
 
 ## mockable bindings: start
 ## mockable bindings: end

R/tidymodels.R

@@ -0,0 +1,96 @@
+#' Methods for tidymodels objects
+#'
+#' Fitted \pkg{parsnip} model objects (`model_fit`) and fitted \pkg{workflows}
+#' objects (`workflow`) are unwrapped with [hardhat::extract_fit_engine()] and
+#' the generic is then called on the underlying engine model. All other
+#' arguments, such as `tree`, `add_statistics`, or `data`, are passed to the
+#' method for the engine model; see the documentation for that model's method.
+#'
+#' For `as.party()`, engines that need the training data (e.g., xgboost) require
+#' it via `data`. For workflows, this should be the preprocessed predictors that
+#' were given to the engine.
+#'
+#' @param x,obj,object A fitted `model_fit` or `workflow` object.
+#' @param ... Other arguments passed to the method for the engine model.
+#'
+#' @return The result of the generic applied to the engine model.
+#'
+#' @examples
+#' if (rlang::is_installed(c("parsnip", "ranger"))) {
+#'   library(parsnip)
+#'
+#'   rf_fit <-
+#'     rand_forest(trees = 5) |>
+#'     set_engine("ranger") |>
+#'     set_mode("regression") |>
+#'     fit(mpg ~ ., data = mtcars)
+#'
+#'   num_ensemble_members(rf_fit)
+#'   extract_rules(rf_fit, tree = 1:2, data = mtcars)
+#' }
+#' @name tidymodels-methods
+NULL
+
+extract_engine <- function(x) {
+  hardhat::extract_fit_engine(x)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+extract_rules.model_fit <- function(x, ...) {
+  extract_rules(extract_engine(x), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+extract_rules.workflow <- function(x, ...) {
+  extract_rules(extract_engine(x), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+active_predictors.model_fit <- function(x, ...) {
+  active_predictors(extract_engine(x), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+active_predictors.workflow <- function(x, ...) {
+  active_predictors(extract_engine(x), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+num_ensemble_members.model_fit <- function(x, ...) {
+  num_ensemble_members(extract_engine(x), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+num_ensemble_members.workflow <- function(x, ...) {
+  num_ensemble_members(extract_engine(x), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+as.party.model_fit <- function(obj, ...) {
+  partykit::as.party(extract_engine(obj), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+as.party.workflow <- function(obj, ...) {
+  partykit::as.party(extract_engine(obj), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+var_imp.model_fit <- function(object, ...) {
+  generics::var_imp(extract_engine(object), ...)
+}
+
+#' @rdname tidymodels-methods
+#' @export
+var_imp.workflow <- function(object, ...) {
+  generics::var_imp(extract_engine(object), ...)
+}

inst/WORDLIST

@@ -36,4 +36,6 @@ randomForest
 rightDaughter
 rpart
 tibble
+tidymodels
 xgb
+xgboost

plans/2026-10-26-1717-tidymodels-wrapper-methods.md

@@ -0,0 +1,45 @@
+# Tidymodels wrapper methods (`model_fit`, `workflow`)
+
+## Overview
+
+Context: lorax defines five generics with per-engine S3 methods: `extract_rules()`, `active_predictors()`, `num_ensemble_members()` (lorax), `as.party()` (partykit, first arg `obj`), and `var_imp()` (generics, first arg `object`). Users with parsnip `model_fit` or `workflow` objects must currently call `extract_fit_engine()` by hand. Add methods for both classes on all five generics that unwrap with `generics::extract_fit_engine()` and re-dispatch on the engine object.
+
+Design: use two small internal helpers and one-line methods that forward `...`. Every engine-specific argument (`tree`, `chain`, `committee`, `add_statistics`, `data`, `nthread`, `complete`, ...) already flows through `...`, so no argument duplication is needed and engine methods keep owning validation. Only the first-argument name differs per generic (`x`, `obj`, `object`).
+
+Notes:
+- `generics` is already in Imports, so `extract_fit_engine` is available without a new dependency. `parsnip` is already in Suggests; add `workflows` to Suggests (tests only).
+- `workflow` objects must be fitted; `extract_fit_engine.workflow()` errors otherwise, which is the desired behavior.
+- `as.party()` for xgboost/lightgbm/catboost/C5.0/etc. requires `data`; the user supplies it via `...` (for workflows this is the baked/processed data, documented in the roxygen note).
+- Single-tree engines still have no `num_ensemble_members` method, so the unwrapped dispatch gives the normal "no applicable method" error.
+
+## Checklist
+
+- [ ] Copy this plan to `plans/2026-10-26-{hhmm}-tidymodels-wrapper-methods.md` in the repo (per CLAUDE.md naming; plan mode blocked writing there)
+- [ ] Create `R/tidymodels.R` with helper `extract_engine(x)` (calls `generics::extract_fit_engine(x)`, no inline conditionals) and `@export` methods: `extract_rules.model_fit`, `extract_rules.workflow`, `active_predictors.model_fit`, `active_predictors.workflow`, `num_ensemble_members.model_fit`, `num_ensemble_members.workflow`, `as.party.model_fit`, `as.party.workflow`, `var_imp.model_fit`, `var_imp.workflow`
+- [ ] Each method body: `UseGeneric(extract_engine(x), ...)`, e.g. `extract_rules.model_fit <- function(x, ...) extract_rules(extract_engine(x), ...)`; `as.party` uses `obj`, `var_imp` uses `object`
+- [ ] Roxygen: one shared `@rdname` topic (e.g. `tidymodels-methods`) with `@inheritParams`, a short description, and an example guarded by `rlang::is_installed(c("parsnip", "rpart"))`; add topic to `_pkgdown.yml`
+- [ ] DESCRIPTION: add `workflows` to Suggests
+- [ ] `devtools::document()`; confirm NAMESPACE gains the 10 `S3method()` lines
+- [ ] Tests in `tests/testthat/test-tidymodels.R`: for a parsnip `decision_tree`/`rand_forest` (ranger) `model_fit` and a fitted `workflow`, check each generic equals the result on `extract_fit_engine(fit)`, that `tree`/`add_statistics` pass through, and use `expect_snapshot(error = TRUE)` for an unfitted workflow; skip if parsnip/workflows/engines not installed
+- [ ] NEWS.md bullet (no function name → place at top, unwrapped line)
+- [ ] `air format .`, `pkgdown::check_pkgdown()`
+- [ ] Run `devtools::test(filter = "tidymodels")` then full `devtools::test()`
+- [ ] Leave all work uncommitted (no commit/push/PR without explicit request)
+
+## Critical files
+
+- `R/generics.R` (generic definitions, signature reference)
+- `R/tidymodels.R` (new)
+- `tests/testthat/test-tidymodels.R` (new)
+- `DESCRIPTION`, `NEWS.md`, `_pkgdown.yml`, `NAMESPACE` (generated)
+
+## Verification
+
+```
+Rscript -e "devtools::document()"
+Rscript -e "devtools::test(filter = 'tidymodels')"
+Rscript -e "devtools::test()"
+Rscript -e "pkgdown::check_pkgdown()"
+```
+
+Manual: fit `parsnip::rand_forest() |> set_engine("ranger") |> set_mode("regression") |> fit(mpg ~ ., mtcars)`, wrap in a workflow, and call `extract_rules(x, tree = 2)`, `active_predictors()`, `num_ensemble_members()`, `var_imp()` on both.

tests/testthat/_snaps/tidymodels.md

@@ -0,0 +1,9 @@
+# unfitted workflows error
+
+    Code
+      extract_rules(unfitted)
+    Condition
+      Error in `extract_fit_parsnip()`:
+      ! Can't extract a model fit from an untrained workflow.
+      i Do you need to call `fit()`?
+

tests/testthat/test-tidymodels.R

@@ -0,0 +1,68 @@
+skip_if_not_installed("parsnip")
+skip_if_not_installed("workflows")
+skip_if_not_installed("ranger")
+
+rf_spec <-
+  parsnip::rand_forest(trees = 5) |>
+  parsnip::set_engine("ranger", importance = "impurity") |>
+  parsnip::set_mode("regression")
+
+rf_fit <- parsnip::fit(rf_spec, mpg ~ ., data = mtcars)
+rf_wflow <-
+  workflows::workflow(mpg ~ ., rf_spec) |>
+  parsnip::fit(data = mtcars)
+rf_engine <- parsnip::extract_fit_engine(rf_fit)
+wflow_engine <- parsnip::extract_fit_engine(rf_wflow)
+
+test_that("extract_rules() unwraps model_fit and workflow objects", {
+  expect_equal(
+    extract_rules(rf_fit, tree = 2:3, data = mtcars, add_statistics = TRUE),
+    extract_rules(rf_engine, tree = 2:3, data = mtcars, add_statistics = TRUE)
+  )
+  expect_equal(
+    extract_rules(rf_wflow, tree = 2:3, data = mtcars, add_statistics = TRUE),
+    extract_rules(
+      wflow_engine,
+      tree = 2:3,
+      data = mtcars,
+      add_statistics = TRUE
+    )
+  )
+})
+
+test_that("active_predictors() unwraps model_fit and workflow objects", {
+  expect_equal(
+    active_predictors(rf_fit, tree = 1:2),
+    active_predictors(rf_engine, tree = 1:2)
+  )
+  expect_equal(
+    active_predictors(rf_wflow, tree = 1:2),
+    active_predictors(wflow_engine, tree = 1:2)
+  )
+})
+
+test_that("num_ensemble_members() unwraps model_fit and workflow objects", {
+  expect_equal(num_ensemble_members(rf_fit), 5L)
+  expect_equal(num_ensemble_members(rf_wflow), 5L)
+})
+
+test_that("var_imp() unwraps model_fit and workflow objects", {
+  expect_equal(var_imp(rf_fit), var_imp(rf_engine))
+  expect_equal(var_imp(rf_wflow), var_imp(wflow_engine))
+})
+
+test_that("as.party() unwraps model_fit and workflow objects", {
+  expect_equal(
+    extract_rules(partykit::as.party(rf_fit, tree = 2, data = mtcars)),
+    extract_rules(partykit::as.party(rf_engine, tree = 2, data = mtcars))
+  )
+  expect_equal(
+    extract_rules(partykit::as.party(rf_wflow, tree = 2, data = mtcars)),
+    extract_rules(partykit::as.party(wflow_engine, tree = 2, data = mtcars))
+  )
+})
+
+test_that("unfitted workflows error", {
+  unfitted <- workflows::workflow(mpg ~ ., rf_spec)
+  expect_snapshot(extract_rules(unfitted), error = TRUE)
+})