feat(users/Profpatsch/my-prelude): add RunCommand.hs r/6103

Change-Id: I08231027a7363ba89006e4dcd510302599be7b4c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8499
Tested-by: BuildkiteCI
Autosubmit: Profpatsch <mail@profpatsch.de>
Reviewed-by: Profpatsch <mail@profpatsch.de>

5 files changed, 174 insertions, 0 deletions

diff --git a/users/Profpatsch/my-prelude/Pretty.hs b/users/Profpatsch/my-prelude/Pretty.hs
index 8a58a5934e..8046c83e45 100644
--- a/users/Profpatsch/my-prelude/Pretty.hs
+++ b/users/Profpatsch/my-prelude/Pretty.hs
@@ -5,6 +5,7 @@ module Pretty
   ( -- * Pretty printing for error messages
     Err,
     printPretty,
+    showPretty,
     -- constructors hidden
     prettyErrs,
     message,
@@ -40,6 +41,9 @@ printPretty :: Show a => a -> IO ()
 printPretty a =
   a & pretty & (: []) & prettyErrs & stringToText & putStderrLn
 
+showPretty :: Show a => a -> Text
+showPretty a = a & pretty & (: []) & prettyErrs & stringToText
+
 -- | Display a list of 'Err's as a colored error message
 -- and abort the test.
 prettyErrs :: [Err] -> String
diff --git a/users/Profpatsch/my-prelude/README.md b/users/Profpatsch/my-prelude/README.md
index 8173d76680..2cc068579a 100644
--- a/users/Profpatsch/my-prelude/README.md
+++ b/users/Profpatsch/my-prelude/README.md
@@ -36,3 +36,7 @@ Helpers around Json parsing.
 ### `Data.Error.Tree`
 
 Collect errors (from [`Data.Error`](https://hackage.haskell.org/package/error-1.0.0.0/docs/Data-Error.html)) into a tree, then display them in a nested fashion. Super useful for e.g. collecting and displaying nested parsing errors.
+
+### `RunCommand.hs`
+
+A module wrapping the process API with some helpful defaults for executing commands and printing what is executed to stderr.
diff --git a/users/Profpatsch/my-prelude/RunCommand.hs b/users/Profpatsch/my-prelude/RunCommand.hs
new file mode 100644
index 0000000000..5c80eb3aac
--- /dev/null
+++ b/users/Profpatsch/my-prelude/RunCommand.hs
@@ -0,0 +1,162 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE QuasiQuotes #-}
+
+module RunCommand where
+
+import Data.ByteString qualified as ByteString
+import Data.ByteString.Lazy qualified as Bytes.Lazy
+import Data.Char qualified as Char
+import Data.List qualified as List
+import Data.Text qualified as Text
+import MyPrelude
+import System.Exit qualified as Exit
+import System.IO (Handle)
+import System.Process.Typed qualified as Process
+import Prelude hiding (log)
+
+-- | Given a a command, the executable and arguments,
+-- spawn the tool as subprocess and collect its stdout (stderr will go to our stderr).
+
+-- Will strip the stdout of trailing newlines.
+--
+-- If the executable is not a path, it will be resolved via the @PATH@ environment variable.
+runCommand :: MonadIO m => FilePath -> [Text] -> m (Exit.ExitCode, ByteString)
+runCommand executable args = do
+  let bashArgs = prettyArgsForBash ((executable & stringToText) : args)
+  log [fmt|Running: $ {bashArgs}|]
+  Process.proc
+    executable
+    (args <&> textToString)
+    & Process.readProcessStdout
+    <&> second toStrictBytes
+    <&> second stripWhitespaceFromEnd
+
+-- | Given a a command, the executable and arguments,
+-- spawn the tool as subprocess and run it to conclusion.
+--
+-- If the executable is not a path, it will be resolved via the @PATH@ environment variable.
+runCommandNoStdout :: MonadIO m => FilePath -> [Text] -> m Exit.ExitCode
+runCommandNoStdout executable args = do
+  let bashArgs = prettyArgsForBash ((executable & stringToText) : args)
+  log [fmt|Running: $ {bashArgs}|]
+  Process.proc
+    executable
+    (args <&> textToString)
+    & Process.runProcess
+
+-- TODO: This is reversing the whole string *twice*. Can we strip from end without doing that?
+stripWhitespaceFromEnd :: ByteString -> ByteString
+stripWhitespaceFromEnd = ByteString.reverse . ByteString.dropWhile (\w -> w == charToWordUnsafe '\n') . ByteString.reverse
+
+-- | Like `runCommand`, but takes a Bytestring that provides the command with streamed input on stdin.
+runCommandWithStdin :: MonadIO m => FilePath -> [Text] -> Bytes.Lazy.ByteString -> m (Exit.ExitCode, ByteString)
+runCommandWithStdin executable args stdin = do
+  let bashArgs = prettyArgsForBash ((executable & stringToText) : args)
+  log [fmt|Running: $ {bashArgs}|]
+  Process.proc
+    executable
+    (args <&> textToString)
+    & Process.setStdin (Process.byteStringInput stdin)
+    & Process.readProcessStdout
+    <&> second toStrictBytes
+    <&> second stripWhitespaceFromEnd
+
+-- | Like `runCommand`, but takes a Bytestring that provides the command with streamed input on stdin.
+runCommandWithStdinNoStdout :: MonadIO m => FilePath -> [Text] -> Bytes.Lazy.ByteString -> m Exit.ExitCode
+runCommandWithStdinNoStdout executable args stdin = do
+  let bashArgs = prettyArgsForBash ((executable & stringToText) : args)
+  log [fmt|Running: $ {bashArgs}|]
+  Process.proc
+    executable
+    (args <&> textToString)
+    & Process.setStdin (Process.byteStringInput stdin)
+    & Process.runProcess
+
+-- | Like 'runCommandWithStdin' but exit if the command returns a non-0 status.
+runCommandWithStdinExpect0 :: MonadIO m => FilePath -> [Text] -> Bytes.Lazy.ByteString -> m ByteString
+runCommandWithStdinExpect0 executable args stdin =
+  runCommandWithStdin executable args stdin >>= \case
+    (ex, stdout) -> do
+      checkStatus0 executable ex
+      pure stdout
+
+-- | Like 'runCommandWithStdinNoStdout' but exit if the command returns a non-0 status.
+runCommandWithStdinNoStdoutExpect0 :: MonadIO m => FilePath -> [Text] -> Bytes.Lazy.ByteString -> m ()
+runCommandWithStdinNoStdoutExpect0 executable args stdin =
+  runCommandWithStdinNoStdout executable args stdin
+    >>= checkStatus0 executable
+
+-- | Like 'runCommandExpect0', but don’t capture stdout,
+-- connect stdin and stdout to the command until it returns.
+--
+-- This is for interactive subcommands.
+runCommandInteractiveExpect0 :: MonadIO m => FilePath -> [Text] -> m ()
+runCommandInteractiveExpect0 executable args = do
+  let bashArgs = prettyArgsForBash ((executable & stringToText) : args)
+  log [fmt|Running interactively: $ {bashArgs}|]
+  ( liftIO $
+      Process.runProcess $
+        Process.proc
+          executable
+          (args <&> textToString)
+    )
+    >>= checkStatus0 executable
+
+-- | Given a name of a command, the executable and arguments,
+-- spawn the tool as subprocess and pipe its stdout to the given 'Handle'.
+--
+-- If the executable is not a path, it will be resolved via the @PATH@ environment variable.
+runCommandPipeToHandle :: MonadIO m => FilePath -> [Text] -> Handle -> m Exit.ExitCode
+runCommandPipeToHandle executable args handle = do
+  -- TODO log the output file?
+  let bashArgs = prettyArgsForBash ((executable & stringToText) : args)
+  log [fmt|Running: $ {bashArgs}|]
+  liftIO $
+    Process.runProcess
+      ( Process.proc
+          executable
+          (args <&> textToString)
+          & Process.setStdout (Process.useHandleClose handle)
+      )
+
+-- | Check whether a command exited 0 or crash.
+checkStatus0 :: MonadIO m => FilePath -> Exit.ExitCode -> m ()
+checkStatus0 executable = \case
+  Exit.ExitSuccess -> pure ()
+  Exit.ExitFailure status -> do
+    logCritical [fmt|Command `{executable}` did not exit with status 0 (success), but status {status}|]
+
+log :: MonadIO m => Text -> m ()
+log = liftIO . putStderrLn
+
+-- | Log the message on the normal logging level & exit the program
+logCritical :: MonadIO m => Text -> m b
+logCritical msg = do
+  liftIO $ putStderrLn msg
+  liftIO $ Exit.exitWith (Exit.ExitFailure 1)
+
+-- | Pretty print a command line in a way that can be copied to bash.
+prettyArgsForBash :: [Text] -> Text
+prettyArgsForBash = Text.intercalate " " . map simpleBashEscape
+
+-- | Simple escaping for bash words. If they contain anything that’s not ascii chars
+-- and a bunch of often-used special characters, put the word in single quotes.
+simpleBashEscape :: Text -> Text
+simpleBashEscape t = do
+  case Text.find (not . isSimple) t of
+    Just _ -> escapeSingleQuote t
+    Nothing -> t
+  where
+    -- any word that is just ascii characters is simple (no spaces or control characters)
+    -- or contains a few often-used characters like - or .
+    isSimple c =
+      Char.isAsciiLower c
+        || Char.isAsciiUpper c
+        || Char.isDigit c
+        -- These are benign, bash will not interpret them as special characters.
+        || List.elem c ['-', '.', ':', '/']
+    -- Put the word in single quotes
+    -- If there is a single quote in the word,
+    -- close the single quoted word, add a single quote, open the word again
+    escapeSingleQuote t' = "'" <> Text.replace "'" "'\\''" t' <> "'"
diff --git a/users/Profpatsch/my-prelude/default.nix b/users/Profpatsch/my-prelude/default.nix
index e4df48dfc9..4d854b194b 100644
--- a/users/Profpatsch/my-prelude/default.nix
+++ b/users/Profpatsch/my-prelude/default.nix
@@ -11,6 +11,7 @@ pkgs.haskellPackages.mkDerivation {
     ./Pretty.hs
     ./Data/Error/Tree.hs
     ./Aeson.hs
+    ./RunCommand.hs
     ./Test.hs
   ];
 
@@ -30,6 +31,7 @@ pkgs.haskellPackages.mkDerivation {
     pkgs.haskellPackages.hspec-expectations-pretty-diff
     pkgs.haskellPackages.hscolour
     pkgs.haskellPackages.nicify-lib
+    pkgs.haskellPackages.typed-process
     pkgs.haskellPackages.ansi-terminal
     pkgs.haskellPackages.vector
   ];
diff --git a/users/Profpatsch/my-prelude/my-prelude.cabal b/users/Profpatsch/my-prelude/my-prelude.cabal
index fd02578013..94f9c9e774 100644
--- a/users/Profpatsch/my-prelude/my-prelude.cabal
+++ b/users/Profpatsch/my-prelude/my-prelude.cabal
@@ -11,6 +11,7 @@ library
       Pretty
       Data.Error.Tree
       Aeson
+      RunCommand
       Test
 
     -- Modules included in this executable, other than Main.
@@ -37,6 +38,7 @@ library
      , hspec-expectations-pretty-diff
      , hscolour
      , nicify-lib
+     , typed-process
      , ansi-terminal
      , vector
     default-language: GHC2021