From 9ae7831f2d0054a721e62451b9deccf3a2a21176 Mon Sep 17 00:00:00 2001
From: Simon Michael <simon@joyful.com>
Date: Wed, 5 Feb 2014 13:56:49 -0800
Subject: [PATCH] extra: hledger-rewrite improvements

---
 extra/hledger-rewrite.hs | 48 ++++++++++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/extra/hledger-rewrite.hs b/extra/hledger-rewrite.hs
index 479bb31c7..1158e1032 100755
--- a/extra/hledger-rewrite.hs
+++ b/extra/hledger-rewrite.hs
@@ -1,11 +1,13 @@
 #!/usr/bin/env runhaskell
 {-|
-hledger-rewrite PATTERNS --add-posting "ACCT  AMTEXPR" ...
+hledger-rewrite [PATTERNS] [--add-posting "ACCT  AMTEXPR"] ...
 
 Skeleton for a minimal generic rewriter of journal entries.
 Reads the default journal and prints the entries, like print,
 but adds the specified postings to any entries matching PATTERNS.
 
+Tested-with: hledger 0.22.2
+
 |-}
 
 import Hledger.Data.Types (Journal(..))
@@ -13,18 +15,32 @@ import Hledger
 import Hledger.Cli
 
 main = do
-  putStrLn "(-f option not supported)"
-  opts <- getCliOpts (defCommandMode ["hledger-rewrite"])
-  withJournalDo opts $
-    \opts j@Journal{jtxns=ts} -> print' opts j{jtxns=map rewrite ts}
-    where 
-      rewrite t = if matched t then t{tpostings=tpostings t ++ newps t}
-                               else t
-      matched t = Acct "^income" `matchesTransaction` t
-      newps t = [generatePosting t "(Reserve)" (`divideMixedAmount` 10)]
-      generatePosting t acct amtfn = nullposting{paccount = accountNameWithoutPostingType acct
-                                                ,ptype    = accountNamePostingType acct
-                                                ,pamount  = amtfn amt
-                                                }
-        where
-          amt = pamount $ head $ filter (Acct "^income" `matchesPosting`) $ tpostings t
+  opts@CliOpts{reportopts_=ropts} <- getCliOpts (defCommandMode ["hledger-rewrite"])
+  d <- getCurrentDay
+  let
+    q = queryFromOpts d ropts
+    -- parse added postings from args.. hard-coded here:
+    addps :: [(AccountName, Transaction -> MixedAmount)]
+    addps = [
+      ("(Reserve)", (\t -> (t `firstAmountMatching` q) `divideMixedAmount` 10))
+      ]
+
+  withJournalDo opts $ \opts j@Journal{jtxns=ts} -> do
+      -- rewrite matched transactions
+      let j' = j{jtxns=map (\t -> if q `matchesTransaction` t then rewriteTransaction t addps else t) ts}
+      -- print' opts j'
+      -- print all transactions (without filtering)
+      putStr $ showTransactions ropts Any j'
+      
+rewriteTransaction :: Transaction -> [(AccountName, Transaction -> MixedAmount)] -> Transaction
+rewriteTransaction t addps = t{tpostings=tpostings t ++ map (uncurry (generatePosting t)) addps}
+
+generatePosting :: Transaction -> AccountName -> (Transaction -> MixedAmount) -> Posting
+generatePosting t acct amtfn = nullposting{paccount     = accountNameWithoutPostingType acct
+                                          ,ptype        = accountNamePostingType acct
+                                          ,pamount      = amtfn t
+                                          ,ptransaction = Just t
+                                          }
+
+firstAmountMatching :: Transaction -> Query -> MixedAmount
+firstAmountMatching t q = pamount $ head $ filter (q `matchesPosting`) $ tpostings t