extra: hledger-rewrite now takes arguments and is usable

extra/hledger-rewrite.hs

@@ -1,45 +1,96 @@
 #!/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.
+A start at a 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
+Examples:
+
+hledger-rewrite.hs ^income --add-posting '(liabilities:tax)  *.33' --add-posting '(reserve:gifts)  $100'
+hledger-rewrite.hs expenses:gifts --add-posting '(reserve:gifts)  *-1"'
+
+Note the single quotes to protect the dollar sign from bash, and the two spaces between account and amount.
+See the command-line help for more details.
+Currently does not work when invoked via "hledger rewrite".
+
+Tested-with: hledger HEAD ~ 2014/2/4
 
 |-}
 
-import Hledger
+-- hledger lib, cli and cmdargs utils
 import Hledger.Cli
+-- more utils for parsing
+import Control.Applicative hiding (many)
+import Text.ParserCombinators.Parsec
 
-main = do
-  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'
-      
+cmdmode :: Mode RawOpts
+cmdmode = (defCommandMode ["hledger-rewrite"]) {
+   modeArgs = ([], Just $ argsFlag "[PATTERNS] --add-posting \"ACCT  AMTEXPR\" ...")
+  ,modeHelp = "show all journal entries, adding specified custom postings to matched ones"
+  ,modeGroupFlags = Group {
+     groupNamed = [("Input",     inputflags)
+                  ,("Reporting", reportflags)
+                  ,("Misc",      helpflags)
+                 ]
+    ,groupUnnamed = [flagReq ["add-posting"] (\s opts -> Right $ setopt "add-posting" s opts) "\"ACCT  AMTEXPR\"" 
+                     "add a posting to ACCT (can be parenthesised) with amount generated by an expression, which is: a literal amount, or * followed by a decimal multiplier (which multiplies the entry's first amount matched by PATTERNS). Two spaces are required between account and amount."]
+    ,groupHidden = []
+    }
+  }
+
+type PostingExpr = (AccountName, AmountExpr)
+
+data AmountExpr = AmountLiteral String | AmountMultiplier Quantity deriving (Show)
+
+addPostingExprsFromOpts :: RawOpts -> [PostingExpr]
+addPostingExprsFromOpts = map (either parseerror id . parseWithCtx nullctx postingexprp) . map stripquotes . listofstringopt "add-posting"
+
+postingexprp = do
+  a <- accountnamep
+  spacenonewline >> many1 spacenonewline
+  aex <- amountexprp
+  many spacenonewline
+  return (a,aex)
+
+amountexprp =
+  choice [
+     AmountMultiplier <$> (do char '*'
+                              many spacenonewline
+                              (q,_,_,_,_) <- numberp
+                              return q)
+    ,AmountLiteral <$> many anyChar
+    ]
+
+amountExprRenderer :: Query -> AmountExpr -> (Transaction -> MixedAmount)
+amountExprRenderer q aex =
+  case aex of 
+    AmountLiteral s    -> either parseerror (const . mixed) $ parseWithCtx nullctx amountp s
+    AmountMultiplier n -> (`divideMixedAmount` (1 / n)) . (`firstAmountMatching` q)
+  where
+    firstAmountMatching :: Transaction -> Query -> MixedAmount
+    firstAmountMatching t q = pamount $ head $ filter (q `matchesPosting`) $ tpostings t
+
 rewriteTransaction :: Transaction -> [(AccountName, Transaction -> MixedAmount)] -> Transaction
 rewriteTransaction t addps = t{tpostings=tpostings t ++ map (uncurry (generatePosting t)) addps}
+  where
+    generatePosting :: Transaction -> AccountName -> (Transaction -> MixedAmount) -> Posting
+    generatePosting t acct amtfn = nullposting{paccount     = accountNameWithoutPostingType acct
+                                              ,ptype        = accountNamePostingType acct
+                                              ,pamount      = amtfn t
+                                              ,ptransaction = Just t
+                                              }
 
-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
+main = do
+  opts@CliOpts{rawopts_=rawopts,reportopts_=ropts} <- getCliOpts cmdmode
+  d <- getCurrentDay
+  let q = queryFromOpts d ropts
+      addps = [(a, amountExprRenderer q aex) | (a, aex) <- addPostingExprsFromOpts rawopts]
+  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}
+    -- run the print command, showing all transactions
+    print' opts{reportopts_=ropts{query_=""}} j'
+