feat: generalize withSetOptionIn to arbitrary result types

src/Lean/Linter/Basic.lean

@@ -23,9 +23,11 @@ namespace Lean
 open Elab Command
 
 /--
-Given a command elaborator `cmd`, returns a new command elaborator that
+Given a function `cmd` on command syntax, returns a new function that
 first peels off and evaluates `set_option ... in ...` syntax recursively, updating the options
-accordingly, and then invokes `cmd` on the inner syntax.
+accordingly, and then invokes `cmd` on the inner syntax. The result type is arbitrary, so that
+besides `CommandElab`s this can also wrap functions that return values, such as the phases of a
+stateful linter.
 
 This is expected to be used in linters, after elaboration is complete. It is not appropriate for
 ordinary elaboration of outer `set_option`s, since it
@@ -34,7 +36,8 @@ ordinary elaboration of outer `set_option`s, since it
   unknown or the wrong type of value is provided), as these should have been reported during
   original elaboration.
 -/
-partial def withSetOptionIn (cmd : CommandElab) : CommandElab := fun stx => do
+partial def withSetOptionIn (cmd : Syntax → CommandElabM α) : Syntax → CommandElabM α :=
+  fun stx => do
   if stx.getKind == ``Lean.Parser.Command.in &&
      stx[0].getKind == ``Lean.Parser.Command.set_option then
       -- Do not modify the infotrees when elaborating, and silently ignore errors.

tests/elab/withSetOptionIn.lean

@@ -132,3 +132,24 @@ set_option linter.all 3 in
 example := trivial
 
 end malformedOption
+
+section resultType
+
+/-!
+This test checks that `withSetOptionIn` can wrap functions with an arbitrary result type,
+such as the phases of a stateful linter. It also checks that the option values are visible
+inside the wrapped function only.
+-/
+
+run_cmd do
+  let stx ← `(command| set_option pp.raw true in set_option pp.explicit true in #check True)
+  let (raw, explicit, inner) ← withSetOptionIn
+    (fun inner => return (pp.raw.get (← getOptions), pp.explicit.get (← getOptions), inner))
+    stx
+  unless raw do throwError "pp.raw was not applied"
+  unless explicit do throwError "pp.explicit was not applied"
+  unless inner.isOfKind ``Lean.Parser.Command.check do
+    throwError "unexpected inner syntax kind: {inner.getKind}"
+  if pp.raw.get (← getOptions) then throwError "options leaked out of `withSetOptionIn`"
+
+end resultType