{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE QuasiQuotes #-} module Hledger.Flow.Common where import Path (absfile, relfile) import qualified Path.IO as Path import qualified Turtle import Turtle ((%), (), (<.>)) import Prelude hiding (putStrLn) import qualified Data.Text as T import qualified Data.Text.IO as T import qualified Data.Text.Read as T import qualified GHC.IO.Handle.FD as H import Data.Char (isDigit) import Data.Either import qualified Control.Foldl as Fold import qualified Data.Map.Strict as Map import Data.Function (on) import qualified Data.List as List (null, sortBy, groupBy) import Data.Ord (comparing) import Hledger.Flow.Types import Hledger.Flow.Logging import Hledger.Flow.PathHelpers (AbsFile, TurtlePath, fromTurtleAbsFile, pathToTurtle) import Hledger.Flow.BaseDir (turtleBaseDir, relativeToBase) import Control.Concurrent.STM hledgerPathFromOption :: Maybe TurtlePath -> IO AbsFile hledgerPathFromOption pathOption = do case pathOption of Just h -> do hlAbs <- fromTurtleAbsFile h isOnDisk <- Path.doesFileExist hlAbs if isOnDisk then return hlAbs else do let msg = Turtle.format ("Unable to find hledger at "%Turtle.fp) h errExit' 1 (T.hPutStrLn H.stderr) msg hlAbs Nothing -> do maybeH <- Path.findExecutable [relfile|hledger|] case maybeH of Just h -> return h Nothing -> do let msg = "Unable to find hledger in your path.\n" <> "You need to either install hledger, or add it to your PATH, or provide the path to an hledger executable.\n\n" <> "There are a number of installation options on the hledger website: https://hledger.org/download.html" errExit' 1 (T.hPutStrLn H.stderr) msg [absfile|/hledger|] hledgerVersionFromPath :: TurtlePath -> IO T.Text hledgerVersionFromPath hlp = fmap (T.strip . Turtle.linesToText) (Turtle.single $ shellToList $ Turtle.inproc (Turtle.format Turtle.fp hlp) ["--version"] Turtle.empty) hledgerInfoFromPath :: Maybe TurtlePath -> IO HledgerInfo hledgerInfoFromPath pathOption = do hlp <- hledgerPathFromOption pathOption hlv <- hledgerVersionFromPath $ pathToTurtle hlp return $ HledgerInfo hlp hlv showCmdArgs :: [T.Text] -> T.Text showCmdArgs args = T.intercalate " " (map escapeArg args) escapeArg :: T.Text -> T.Text escapeArg a = if T.count " " a > 0 then "'" <> a <> "'" else a errExit :: Int -> TChan LogMessage -> T.Text -> a -> IO a errExit exitStatus ch = errExit' exitStatus (channelErrLn ch) errExit' :: Int -> (T.Text -> IO ()) -> T.Text -> a -> IO a errExit' exitStatus logFun errorMessage dummyReturnValue = do logFun errorMessage Turtle.sleep 0.1 _ <- Turtle.exit $ Turtle.ExitFailure exitStatus return dummyReturnValue descriptiveOutput :: T.Text -> T.Text -> T.Text descriptiveOutput outputLabel outTxt = do if not (T.null outTxt) then Turtle.format (Turtle.s%":\n"%Turtle.s%"\n") outputLabel outTxt else "" logTimedAction :: HasVerbosity o => o -> TChan LogMessage -> T.Text -> [T.Text] -> (TChan LogMessage -> T.Text -> IO ()) -> (TChan LogMessage -> T.Text -> IO ()) -> IO FullOutput -> IO FullTimedOutput logTimedAction opts ch cmdLabel extraCmdLabels stdoutLogger stderrLogger action = do logVerbose opts ch $ Turtle.format ("Begin: "%Turtle.s) cmdLabel if (List.null extraCmdLabels) then return () else logVerbose opts ch $ T.intercalate "\n" extraCmdLabels timed@((ec, stdOut, stdErr), diff) <- Turtle.time action stdoutLogger ch stdOut stderrLogger ch stdErr logVerbose opts ch $ Turtle.format ("End: "%Turtle.s%" "%Turtle.s%" ("%Turtle.s%")") cmdLabel (Turtle.repr ec) (Turtle.repr diff) return timed timeAndExitOnErr :: (HasSequential o, HasVerbosity o) => o -> TChan LogMessage -> T.Text -> (TChan LogMessage -> T.Text -> IO ()) -> (TChan LogMessage -> T.Text -> IO ()) -> ProcFun -> ProcInput -> IO FullTimedOutput timeAndExitOnErr opts ch cmdLabel = timeAndExitOnErr' opts ch cmdLabel [] timeAndExitOnErr' :: (HasSequential o, HasVerbosity o) => o -> TChan LogMessage -> T.Text -> [T.Text] -> (TChan LogMessage -> T.Text -> IO ()) -> (TChan LogMessage -> T.Text -> IO ()) -> ProcFun -> ProcInput -> IO FullTimedOutput timeAndExitOnErr' opts ch cmdLabel extraCmdLabels stdoutLogger stderrLogger procFun (cmd, args, stdInput) = do let action = procFun cmd args stdInput timed@((ec, stdOut, stdErr), _) <- logTimedAction opts ch cmdLabel extraCmdLabels stdoutLogger stderrLogger action case ec of Turtle.ExitFailure i -> do let cmdText = Turtle.format (Turtle.s%" "%Turtle.s) cmd $ showCmdArgs args let msgOut = descriptiveOutput "Standard output" stdOut let msgErr = descriptiveOutput "Error output" stdErr let exitMsg = Turtle.format ("\n=== Begin Error: "%Turtle.s%" ===\nExternal command:\n"%Turtle.s%"\nExit code "%Turtle.d%"\n" %Turtle.s%Turtle.s%"=== End Error: "%Turtle.s%" ===\n") cmdLabel cmdText i msgOut msgErr cmdLabel errExit i ch exitMsg timed Turtle.ExitSuccess -> return timed procWithEmptyOutput :: ProcFun procWithEmptyOutput cmd args stdinput = do ec <- Turtle.proc cmd args stdinput return (ec, T.empty, T.empty) parAwareProc :: HasSequential o => o -> ProcFun parAwareProc opts = if (sequential opts) then procWithEmptyOutput else Turtle.procStrictWithErr parAwareActions :: (HasSequential o, HasBatchSize o) => o -> [IO a] -> IO [a] parAwareActions opts = if (sequential opts) then sequence else parBatchedActions (batchSize opts) [] parBatchedActions :: Int -> [a] -> [IO a] -> IO [a] parBatchedActions _ done [] = return done parBatchedActions batch done todo = do let doNow = take batch todo let remaining = drop batch todo doneNow <- (Turtle.single . shellToList . Turtle.parallel) doNow parBatchedActions batch (done ++ doneNow) remaining inprocWithErrFun :: (T.Text -> IO ()) -> ProcInput -> Turtle.Shell Turtle.Line inprocWithErrFun errFun (cmd, args, standardInput) = do result <- Turtle.inprocWithErr cmd args standardInput case result of Right ln -> return ln Left ln -> do (Turtle.liftIO . errFun . Turtle.lineToText) ln Turtle.empty verboseTestFile :: (HasVerbosity o, HasBaseDir o) => o -> TChan LogMessage -> TurtlePath -> IO Bool verboseTestFile opts ch p = do fileExists <- Turtle.testfile p let rel = relativeToBase opts p if fileExists then logVerbose opts ch $ Turtle.format ("Found '"%Turtle.fp%"'") rel else logVerbose opts ch $ Turtle.format ("Looked for but did not find '"%Turtle.fp%"'") rel return fileExists groupPairs' :: (Eq a, Ord a) => [(a, b)] -> [(a, [b])] groupPairs' = map (\ll -> (fst . head $ ll, map snd ll)) . List.groupBy ((==) `on` fst) . List.sortBy (comparing fst) groupPairs :: (Eq a, Ord a) => [(a, b)] -> Map.Map a [b] groupPairs = Map.fromList . groupPairs' pairBy :: (a -> b) -> [a] -> [(b, a)] pairBy keyFun = map (\v -> (keyFun v, v)) groupValuesBy :: (Ord k, Ord v) => (v -> k) -> [v] -> Map.Map k [v] groupValuesBy keyFun = groupPairs . pairBy keyFun allYearsFileName :: TurtlePath allYearsFileName = "all-years" <.> "journal" lsDirs :: TurtlePath -> Turtle.Shell TurtlePath lsDirs = onlyDirs . Turtle.ls onlyDirs :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath onlyDirs = excludeHiddenFiles . excludeWeirdPaths . filterPathsByFileStatus Turtle.isDirectory onlyFiles :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath onlyFiles = excludeHiddenFiles . filterPathsByFileStatus Turtle.isRegularFile filterPathsByFileStatus :: (Turtle.FileStatus -> Bool) -> Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath filterPathsByFileStatus filepred files = do files' <- shellToList files filtered <- filterPathsByFileStatus' filepred [] files' Turtle.select filtered filterPathsByFileStatus' :: (Turtle.FileStatus -> Bool) -> [TurtlePath] -> [TurtlePath] -> Turtle.Shell [TurtlePath] filterPathsByFileStatus' _ acc [] = return acc filterPathsByFileStatus' filepred acc (file:files) = do filestat <- Turtle.stat file let filtered = if (filepred filestat) then file:acc else acc filterPathsByFileStatus' filepred filtered files filterPaths :: (TurtlePath -> IO Bool) -> [TurtlePath] -> Turtle.Shell [TurtlePath] filterPaths = filterPaths' [] filterPaths' :: [TurtlePath] -> (TurtlePath -> IO Bool) -> [TurtlePath] -> Turtle.Shell [TurtlePath] filterPaths' acc _ [] = return acc filterPaths' acc filepred (file:files) = do shouldInclude <- Turtle.liftIO $ filepred file let filtered = if shouldInclude then file:acc else acc filterPaths' filtered filepred files excludeHiddenFiles :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath excludeHiddenFiles paths = do p <- paths case (Turtle.match (Turtle.prefix ".") $ Turtle.format Turtle.fp $ Turtle.filename p) of [] -> Turtle.select [p] _ -> Turtle.select [] excludeWeirdPaths :: Turtle.Shell TurtlePath -> Turtle.Shell TurtlePath excludeWeirdPaths = Turtle.findtree (Turtle.suffix $ Turtle.noneOf "_") firstExistingFile :: [TurtlePath] -> IO (Maybe TurtlePath) firstExistingFile files = do case files of [] -> return Nothing file:fs -> do exists <- Turtle.testfile file if exists then return (Just file) else firstExistingFile fs basenameLine :: TurtlePath -> Turtle.Shell Turtle.Line basenameLine path = case (Turtle.textToLine $ Turtle.format Turtle.fp $ Turtle.basename path) of Nothing -> Turtle.die $ Turtle.format ("Unable to determine basename from path: "%Turtle.fp%"\n") path Just bn -> return bn buildFilename :: [Turtle.Line] -> T.Text -> TurtlePath buildFilename identifiers ext = Turtle.fromText (T.intercalate "-" (map Turtle.lineToText identifiers)) <.> ext shellToList :: Turtle.Shell a -> Turtle.Shell [a] shellToList files = Turtle.fold files Fold.list writeFiles :: IO (Map.Map TurtlePath T.Text) -> IO [TurtlePath] writeFiles fileMap = do m <- fileMap writeFiles' m writeFiles' :: Map.Map TurtlePath T.Text -> IO [TurtlePath] writeFiles' fileMap = do writeTextMap fileMap return $ Map.keys fileMap writeTextMap :: Map.Map TurtlePath T.Text -> IO () writeTextMap = Map.foldlWithKey (\a k v -> a <> Turtle.writeTextFile k v) (return ()) changeExtension :: T.Text -> TurtlePath -> TurtlePath changeExtension ext path = (Turtle.dropExtension path) <.> ext changePathAndExtension :: TurtlePath -> T.Text -> TurtlePath -> TurtlePath changePathAndExtension newOutputLocation newExt = (changeOutputPath newOutputLocation) . (changeExtension newExt) changeOutputPath :: TurtlePath -> TurtlePath -> TurtlePath changeOutputPath newOutputLocation srcFile = mconcat $ map changeSrcDir $ Turtle.splitDirectories srcFile where changeSrcDir file = if file == "1-in/" || file == "2-preprocessed/" then newOutputLocation else file listOwners :: HasBaseDir o => o -> Turtle.Shell TurtlePath listOwners opts = fmap Turtle.basename $ lsDirs $ (turtleBaseDir opts) "import" intPath :: Integer -> TurtlePath intPath = Turtle.fromText . (Turtle.format Turtle.d) includeYears :: TChan LogMessage -> TurtlePath -> IO [Integer] includeYears ch includeFile = do txt <- Turtle.readTextFile includeFile case includeYears' txt of Left msg -> do channelErrLn ch msg return [] Right years -> return years includeYears' :: T.Text -> Either T.Text [Integer] includeYears' txt = case partitionEithers (includeYears'' txt) of (errors, []) -> do let msg = Turtle.format ("Unable to extract years from the following text:\n"%Turtle.s%"\nErrors:\n"%Turtle.s) txt (T.intercalate "\n" $ map T.pack errors) Left msg (_, years) -> Right years includeYears'' :: T.Text -> [Either String Integer] includeYears'' txt = map extractDigits (T.lines txt) extractDigits :: T.Text -> Either String Integer extractDigits txt = fmap fst $ (T.decimal . (T.filter isDigit)) txt