From a3cacca71d8aff0892cb392683c8e84fb3394bed Mon Sep 17 00:00:00 2001
From: Stephen Morgan <morgan.stephen99@gmail.com>
Date: Tue, 24 Aug 2021 15:50:32 +1000
Subject: [PATCH] fix: ui: Make sure RegisterScreen (and consequently
 TransactionScreen) only display forecast transactions when the appropriate
 flag is set.

---
 hledger-ui/Hledger/UI/AccountsScreen.hs | 13 ++-----------
 hledger-ui/Hledger/UI/RegisterScreen.hs |  4 ++--
 hledger-ui/Hledger/UI/UIUtils.hs        | 18 ++++++++++++++++++
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/hledger-ui/Hledger/UI/AccountsScreen.hs b/hledger-ui/Hledger/UI/AccountsScreen.hs
index a3a213612..f87c940b2 100644
--- a/hledger-ui/Hledger/UI/AccountsScreen.hs
+++ b/hledger-ui/Hledger/UI/AccountsScreen.hs
@@ -19,7 +19,7 @@ import Control.Monad.IO.Class (liftIO)
 import Data.List
 import Data.Maybe
 import qualified Data.Text as T
-import Data.Time.Calendar (Day, addDays)
+import Data.Time.Calendar (Day)
 import qualified Data.Vector as V
 import Graphics.Vty (Event(..),Key(..),Modifier(..))
 import Lens.Micro.Platform
@@ -77,16 +77,7 @@ asInit d reset ui@UIState{
                         as = map asItemAccountName displayitems
 
     -- Further restrict the query based on the current period and future/forecast mode.
-    rspec' = rspec{_rsQuery=simplifyQuery $ And [_rsQuery rspec, periodq, excludeforecastq (forecast_ $ inputopts_ copts)]}
-      where
-        periodq = Date $ periodAsDateSpan $ period_ ropts
-        -- Except in forecast mode, exclude future/forecast transactions.
-        excludeforecastq (Just _) = Any
-        excludeforecastq Nothing  =  -- not:date:tomorrow- not:tag:generated-transaction
-          And [
-             Not (Date $ DateSpan (Just $ addDays 1 d) Nothing)
-            ,Not generatedTransactionTag
-          ]
+    rspec' = reportSpecSetFutureAndForecast d (forecast_ $ inputopts_ copts) rspec
 
     -- run the report
     (items,_total) = balanceReport rspec' j
diff --git a/hledger-ui/Hledger/UI/RegisterScreen.hs b/hledger-ui/Hledger/UI/RegisterScreen.hs
index 0f1eb65f6..33c84a258 100644
--- a/hledger-ui/Hledger/UI/RegisterScreen.hs
+++ b/hledger-ui/Hledger/UI/RegisterScreen.hs
@@ -56,7 +56,7 @@ rsSetAccount a forceinclusive scr@RegisterScreen{} =
 rsSetAccount _ _ scr = scr
 
 rsInit :: Day -> Bool -> UIState -> UIState
-rsInit d reset ui@UIState{aopts=_uopts@UIOpts{cliopts_=CliOpts{reportspec_=rspec@ReportSpec{_rsReportOpts=ropts}}}, ajournal=j, aScreen=s@RegisterScreen{..}} =
+rsInit d reset ui@UIState{aopts=_uopts@UIOpts{cliopts_=copts@CliOpts{reportspec_=rspec@ReportSpec{_rsReportOpts=ropts}}}, ajournal=j, aScreen=s@RegisterScreen{..}} =
   ui{aScreen=s{rsList=newitems'}}
   where
     -- gather arguments and queries
@@ -72,7 +72,7 @@ rsInit d reset ui@UIState{aopts=_uopts@UIOpts{cliopts_=CliOpts{reportspec_=rspec
         -- always show historical balance
       -- , balanceaccum_= Historical
       }
-    rspec' =
+    rspec' = reportSpecSetFutureAndForecast d (forecast_ $ inputopts_ copts) .
       either (error "rsInit: adjusting the query for register, should not have failed") id $ -- PARTIAL:
       updateReportSpec ropts' rspec{_rsDay=d}
     items = accountTransactionsReport rspec' j thisacctq
diff --git a/hledger-ui/Hledger/UI/UIUtils.hs b/hledger-ui/Hledger/UI/UIUtils.hs
index a0c98f25f..404ccf5de 100644
--- a/hledger-ui/Hledger/UI/UIUtils.hs
+++ b/hledger-ui/Hledger/UI/UIUtils.hs
@@ -24,6 +24,7 @@ module Hledger.UI.UIUtils (
   ,scrollSelectionToMiddle
   ,suspend
   ,redraw
+  ,reportSpecSetFutureAndForecast
 )
 where
 
@@ -36,6 +37,7 @@ import Brick.Widgets.List (List, listSelectedL, listNameL, listItemHeightL)
 import Control.Monad.IO.Class
 import Data.List
 import qualified Data.Text as T
+import Data.Time (Day, addDays)
 import Graphics.Vty
   (Event(..),Key(..),Modifier(..),Vty(..),Color,Attr,currentAttr,refresh
   -- ,Output(displayBounds,mkDisplayContext),DisplayContext(..)
@@ -342,3 +344,19 @@ normaliseMovementKeys ev
   | ev `elem` moveLeftEvents  = EvKey KLeft []
   | ev `elem` moveRightEvents = EvKey KRight []
   | otherwise = ev
+
+-- | Update the ReportSpec's query to exclude future transactions (later than the given day)
+-- and forecast transactions (generated by --forecast), if the given forecast DateSpan is Nothing,
+-- and include them otherwise.
+reportSpecSetFutureAndForecast :: Day -> Maybe DateSpan -> ReportSpec -> ReportSpec
+reportSpecSetFutureAndForecast d forecast rspec =
+    rspec{_rsQuery=simplifyQuery $ And [_rsQuery rspec, periodq, excludeforecastq forecast]}
+  where
+    periodq = Date . periodAsDateSpan . period_ $ _rsReportOpts rspec
+    -- Except in forecast mode, exclude future/forecast transactions.
+    excludeforecastq (Just _) = Any
+    excludeforecastq Nothing  =  -- not:date:tomorrow- not:tag:generated-transaction
+      And [
+         Not (Date $ DateSpan (Just $ addDays 1 d) Nothing)
+        ,Not generatedTransactionTag
+      ]