fix:ui: re-check balance assertions properly when --pivot is used [#2451]

When hledger-ui is started with --pivot, re-enabling balance assertions with the I key now does a full journal reload, to check balance assertions accurately. It means that in pivot mode, the I key can also show other data changes (as if you pressed the g key).
2025-10-09 08:36:12 -10:00 · 2025-10-09 08:36:12 -10:00 · 4aa7d7e20d
commit 4aa7d7e20d
parent 1c86e02d99
5 changed files with 40 additions and 18 deletions
--- a/hledger-ui/Hledger/UI/AccountsScreen.hs
+++ b/hledger-ui/Hledger/UI/AccountsScreen.hs
@ -48,7 +48,7 @@ import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
 import Hledger.UI.Editor
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 import Hledger.UI.RegisterScreen (rsCenterSelection)
 import Data.Either (fromRight)
 import Control.Arrow ((>>>))
@ -257,7 +257,7 @@ asHandleNormalMode (ALS scons ass) ev = do
    VtyEvent (EvKey (KLeft)     [MShift]) -> modify' (previousReportPeriod journalspan >>> regenerateScreens j d)

    -- various toggles and settings:
-    VtyEvent (EvKey (KChar 'I') []) -> modify' (toggleIgnoreBalanceAssertions >>> uiCheckBalanceAssertions d)
+    VtyEvent (EvKey (KChar 'I') []) -> get' >>= uiToggleBalanceAssertions d
    VtyEvent (EvKey (KChar 'F') []) -> modify' (toggleForecast d   >>> regenerateScreens j d)
    VtyEvent (EvKey (KChar 'B') []) -> modify' (toggleConversionOp >>> regenerateScreens j d)
    VtyEvent (EvKey (KChar 'V') []) -> modify' (toggleValue        >>> regenerateScreens j d)
--- a/hledger-ui/Hledger/UI/ErrorScreen.hs
+++ b/hledger-ui/Hledger/UI/ErrorScreen.hs
@ -13,6 +13,7 @@ module Hledger.UI.ErrorScreen
 ,uiCheckBalanceAssertions
 ,uiReload
 ,uiReloadIfFileChanged
+ ,uiToggleBalanceAssertions
 )
 where

@ -102,7 +103,7 @@ esHandle ev = do
                runEditor pos f
              esReloadIfFileChanged copts d j ui

-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (popScreen $ toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d (popScreen ui)
            VtyEvent (EvKey (KChar 'l') [MCtrl]) -> redraw
            VtyEvent (EvKey (KChar 'z') [MCtrl]) -> suspend ui
            _ -> return ()
@ -150,6 +151,10 @@ hledgerparseerrorpositionp = do

 -- Defined here so it can reference the error screen:

+-- | Modify some input options for hledger-ui (enable --forecast).
+uiAdjustOpts :: UIOpts -> CliOpts -> CliOpts
+uiAdjustOpts uopts = enableForecast uopts
+
 -- | Reload the journal from its input files, then update the ui app state accordingly.
 -- This means regenerate the entire screen stack from top level down to the current screen, using the provided today-date.
 -- As a convenience (usually), if journal reloading fails, this enters the error screen, or if already there, updates its message.
@ -163,8 +168,8 @@ hledgerparseerrorpositionp = do
 uiReload :: CliOpts -> Day -> UIState -> EventM Name UIState UIState
 uiReload copts d ui = liftIO $ do
  ej <-
-    let copts' = enableForecast (astartupopts ui) copts
-    in runExceptT $ journalTransform copts' <$> journalReload copts'
+    let copts1 = uiAdjustOpts (astartupopts ui) copts
+    in runExceptT $ journalTransform copts1 <$> journalReload copts1
  -- dbg1IO "uiReload before reload" (map tdescription $ jtxns $ ajournal ui)
  return $ case ej of
    Right j  ->
@ -188,20 +193,24 @@ uiReload copts d ui = liftIO $ do
 -- Also, this one runs in IO, suitable for suspendAndResume.
 uiReloadIfFileChanged :: CliOpts -> Day -> Journal -> UIState -> IO UIState
 uiReloadIfFileChanged copts d j ui = do
-  let copts' = enableForecast (astartupopts ui) copts
-  ej <- runExceptT $ journalReloadIfChanged copts' d j
+  ej <-
+    let copts1 = uiAdjustOpts (astartupopts ui) copts
+    in runExceptT $ journalReloadIfChanged copts1 d j
  return $ case ej of
    Right (j', _) -> regenerateScreens j' d ui
    Left err -> case aScreen ui of
        ES _ -> ui{aScreen=esNew err}
        _    -> pushScreen (esNew err) ui

-- Re-check any balance assertions in the current journal, and if any
-- fail, enter (or update) the error screen. Or if balance assertions
-- are disabled, do nothing.
+-- Re-check any balance assertions in the current journal,
+-- and if any fail, enter (or update) the error screen.
+-- Or if balance assertions are disabled or pivot is active, do nothing.
+-- (When pivot is active, assertions have already been checked on the pre-pivot journal,
+-- and the current post-pivot journal's account names don't match the original assertions.)
 uiCheckBalanceAssertions :: Day -> UIState -> UIState
-uiCheckBalanceAssertions _d ui@UIState{ajournal=j}
-  | ui^.ignore_assertions = ui
+uiCheckBalanceAssertions _d ui@UIState{ajournal=j, aopts=UIOpts{uoCliOpts=CliOpts{inputopts_=InputOpts{pivot_=pval}}}}
+  | ui^.ignore_assertions = ui        -- user disabled checks
+  | not (null pval) = ui              -- post-pivot journal, assertions already checked pre-pivot
  | otherwise =
    case journalCheckBalanceAssertions j of
      Right () -> ui
@ -209,3 +218,16 @@ uiCheckBalanceAssertions _d ui@UIState{ajournal=j}
        case ui of
          UIState{aScreen=ES sst} -> ui{aScreen=ES sst{_essError=err}}
          _                        -> pushScreen (esNew err) ui
+
+-- | Toggle ignoring balance assertions (when user presses I), and if no longer ignoring, recheck them.
+-- Normally the recheck is done quickly on the in-memory journal.
+-- But if --pivot is active, a full journal reload is done instead
+-- (because we can't check balance assertions after pivoting has occurred).
+-- In that case, this operation could be slower and could reveal other data changes (not just balance assertion failures).
+uiToggleBalanceAssertions :: Day -> UIState -> EventM Name UIState ()
+uiToggleBalanceAssertions d ui@UIState{aopts=UIOpts{uoCliOpts=copts@CliOpts{inputopts_=InputOpts{pivot_=pivotval}}}} =
+  let ui' = toggleIgnoreBalanceAssertions ui
+  in case (ui'^.ignore_assertions, null pivotval) of
+    (True, _)      -> put' ui'                                -- ignoring enabled, no check needed
+    (False, True)  -> put' $ uiCheckBalanceAssertions d ui'   -- unpivoted journal, can check in memory
+    (False, False) -> uiReload copts d ui' >>= put'           -- pivoted journal, must reload to check it
--- a/hledger-ui/Hledger/UI/MenuScreen.hs
+++ b/hledger-ui/Hledger/UI/MenuScreen.hs
@ -33,7 +33,7 @@ import Hledger.UI.UITypes
 import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 import Hledger.UI.Editor (runIadd, runEditor, endPosition)
 import Brick.Widgets.Edit (getEditContents, handleEditorEvent)
 import Control.Arrow ((>>>))
@ -156,7 +156,7 @@ msHandle ev = do
              where
                p = reportPeriod ui
            e | e `elem` [VtyEvent (EvKey (KChar 'g') []), AppEvent FileChange] -> uiReload copts d ui >>= put'
-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d ui
            VtyEvent (EvKey (KChar 'a') []) -> suspendAndResume $ clearScreen >> setCursorPosition 0 0 >> add (cliOptsDropArgs copts) j >> uiReloadIfFileChanged copts d j ui
            VtyEvent (EvKey (KChar 'A') []) -> suspendAndResume $ void (runIadd (journalFilePath j)) >> uiReloadIfFileChanged copts d j ui
            VtyEvent (EvKey (KChar 'E') []) -> suspendAndResume $ void (runEditor endPosition (journalFilePath j)) >> uiReloadIfFileChanged copts d j ui
--- a/hledger-ui/Hledger/UI/RegisterScreen.hs
+++ b/hledger-ui/Hledger/UI/RegisterScreen.hs
@ -43,7 +43,7 @@ import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
 import Hledger.UI.Editor
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)

 rsDraw :: UIState -> [Widget Name]
 rsDraw UIState{aopts=_uopts@UIOpts{uoCliOpts=copts@CliOpts{reportspec_=rspec}}
@ -247,7 +247,7 @@ rsHandle ev = do

            e | e `elem` [AppEvent FileChange, VtyEvent (EvKey (KChar 'g') [])] -> uiReload copts d ui >>= put'

-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d ui
            VtyEvent (EvKey (KChar 'a') []) -> suspendAndResume $ clearScreen >> setCursorPosition 0 0 >> add (cliOptsDropArgs copts) j >> uiReloadIfFileChanged copts d j ui
            VtyEvent (EvKey (KChar 'A') []) -> suspendAndResume $ void (runIadd (journalFilePath j)) >> uiReloadIfFileChanged copts d j ui
            VtyEvent (EvKey (KChar 'T') []) -> put' $ regenerateScreens j d $ setReportPeriod (DayPeriod d) ui
--- a/hledger-ui/Hledger/UI/TransactionScreen.hs
+++ b/hledger-ui/Hledger/UI/TransactionScreen.hs
@ -30,7 +30,7 @@ import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
 import Hledger.UI.Editor
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 import Hledger.UI.RegisterScreen (rsHandle)

 tsDraw :: UIState -> [Widget Name]
@ -162,7 +162,7 @@ tsHandle ev = do
              where
                p = reportPeriod ui

-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d ui

            -- for toggles that may change the current/prev/next transactions,
            -- we must regenerate the transaction list, like the g handler above ? with regenerateTransactions ? TODO WIP