Make it workHEADmaster

5 files changed, 88 insertions, 45 deletions

diff --git a/Main.hs b/Main.hs
index d52fedc..30f8f0e 100644
--- a/Main.hs
+++ b/Main.hs
@@ -1,10 +1,12 @@
 {-# LANGUAGE LambdaCase #-}
 module Main where
 
+import Control.Monad (forM_, when)
 import Data.Text.Lazy (Text)
 import qualified Data.Text.Lazy as T
+import qualified Data.Text.Lazy.Builder as B
 import qualified Data.Text.Lazy.IO as T
-import System.Environment
+import Options.Applicative
 import System.Exit
 
 import Pattern
@@ -25,24 +27,47 @@ apply' (Pattern name items) origocc@(Occurrence name' args)
 
 travOrigsArgs :: [Origin] -> [(Text, Text)] -> [(Text, Text)]
 travOrigsArgs [] args = args
-travOrigsArgs _ [] = []
 travOrigsArgs (OCopy : items) (arg:args) = arg : travOrigsArgs items args
 travOrigsArgs (ONew name : items) args = (T.pack " ", name) : travOrigsArgs items args
+travOrigsArgs _ [] = []
+
+data Options = Options { optPattern :: Pattern
+                       , optFiles :: [FilePath]
+                       , optInPlace :: Bool }
+  deriving (Show)
+
+parseOptions :: Parser Options
+parseOptions = Options <$> option (eitherReader parsePattern) (short 'p' <> help "The pattern to apply. *-prefixed type variables are added." <> metavar "PAT")
+                       <*> some (argument str (metavar "FILES..."))
+                       <*> switch (short 'i' <> help "Apply changes in-place. WARNING: This mutates files! \
+                                                     \(This also allows multiple files to be given on the command line.)")
 
 main :: IO ()
 main = do
-    getArgs >>= \case
-        [fname, pat] -> do
-            pat' <- case parsePattern pat of
-              Left err -> die err
-              Right pat' -> return pat'
-            source <- readFile fname
-            parsed <- case parseSourceFile fname source pat' of
-              Left err -> die (show err)
-              Right res -> return res
-            print parsed
-            T.putStrLn (pretty parsed)
-            putStrLn (replicate 80 '-')
-            T.putStrLn (pretty (apply pat' parsed))
-        _ -> do
-            die "Usage: refactor-type-parameters <file.hs> <OpenAcc env aenv *taenv t>"
+    opts <- execParser $
+        info (parseOptions <**> helper)
+             (fullDesc
+              <> progDesc "Refactor type variables of a data type in Haskell"
+              <> header "refactor-type-variables - Refactoring Haskell data type parameters")
+
+    when (not (optInPlace opts) && length (optFiles opts) /= 1) $
+        die "Multiple input files can only be given in conjunction with -i"
+
+    forM_ (optFiles opts) $ \fname -> do
+        source <- readFile fname
+        parsed <- case parseSourceFile fname source (optPattern opts) of
+          Left err -> die ("Could not parse '" ++ fname ++ "': " ++ show err)
+          Right res -> return res
+
+        let result = apply (optPattern opts) parsed
+            resultText = B.toLazyText (sourceFileToBuilder False result)
+
+        if optInPlace opts
+            then do T.writeFile fname resultText
+                    putStrLn ("Processed '" ++ fname ++ "'")
+            else T.putStrLn (pretty (apply (optPattern opts) parsed))
+
+        -- print parsed
+        -- T.putStrLn (pretty parsed)
+        -- putStrLn (replicate 80 '-')
+        -- T.putStrLn (pretty (apply (optPattern opts) parsed))
diff --git a/Parser.hs b/Parser.hs
index afc5b15..017544a 100644
--- a/Parser.hs
+++ b/Parser.hs
@@ -25,18 +25,23 @@ parseSourceFile fname source pat = runReader (runParserT pSourceFile () fname so
 
 pSourceFile :: Parser SourceFile
 pSourceFile = do
-    SourceFile . concat <$> many (pDataDef <|> pInstanceDef <|> pFunDef <|> skipLine)
-  where
-    skipLine :: Parser [Chunk]
-    skipLine =
-        fmap (pure . CSkip . build) $
-            (\s -> B.fromString (s ++ "\n")) <$> manyTill anyChar newline
+    result <- concat <$> many (pDataDef <|> pInstanceDef <|> pFunDef <|> skipLine)
+    after <- skipLineEOF
+    return (SourceFile (result ++ after))
+
+skipLine :: Parser [Chunk]
+skipLine = try $
+    fmap (pure . CSkip . build) $
+         (\s -> B.fromString (s ++ "\n")) <$> manyTill anyChar newline
+
+skipLineEOF :: Parser [Chunk]
+skipLineEOF = pure . CSkip . T.pack <$> manyTill (notFollowedBy newline >> anyChar) eof
 
 -- Parses a function definition on the current line.
 -- Assumes the cursor is at the start of a line, and ends at the start of a line again.
 -- Fails if no function definition can be found on this line.
 pFunDef :: Parser [Chunk]
-pFunDef = try $ do
+pFunDef = flip label "function definition" $ try $ do
     prefix <- build <$> combine
             [horSpaces
             ,option mempty (text "=" +++ horSpaces)
@@ -58,8 +63,8 @@ pFunDef = try $ do
     return (CSkip (prefix `T.append` build s1) : parseTypeText pat typeText ++ [CSkip (build s2)])
 
 pDataDef :: Parser [Chunk]
-pDataDef = try $ do
-    s1 <- text "data"
+pDataDef = flip label "data/type definition" $ try $ do
+    s1 <- text "data" <|> text "type" <|> text "newtype"
     let limiter = lookAhead $ try $
             choice [eof
                    ,do _ <- newline
@@ -71,7 +76,7 @@ pDataDef = try $ do
     return (CSkip (build s1) : parseTypeText pat typeText ++ [CSkip (build s2)])
 
 pInstanceDef :: Parser [Chunk]
-pInstanceDef = try $ do
+pInstanceDef = flip label "instance definition" $ try $ do
     s1 <- text "instance"
     let limiter = lookAhead $ try $
             choice [eof
@@ -80,8 +85,9 @@ pInstanceDef = try $ do
                        _ <- string "where"
                        eof <|> (satisfy prewordc >> return ())]
     typeText <- manyTill anyChar limiter
+    rest <- skipLine
     pat <- ask
-    return (CSkip (build s1) : parseTypeText pat typeText)
+    return (CSkip (build s1) : parseTypeText pat typeText ++ rest)
 
 parseTypeText :: Pattern -> String -> [Chunk]
 parseTypeText (Pattern patname _) inputText =
diff --git a/Pretty.hs b/Pretty.hs
index a5b65a6..2f0afdc 100644
--- a/Pretty.hs
+++ b/Pretty.hs
@@ -3,6 +3,8 @@ module Pretty where
 import qualified Data.Text.Lazy.Builder as B
 import Data.Text.Lazy.Builder (Builder)
 import Data.Text.Lazy (Text)
+import System.IO (hIsTerminalDevice, stdout)
+import System.IO.Unsafe (unsafePerformIO)
 
 
 class Pretty a where
@@ -10,3 +12,6 @@ class Pretty a where
 
     pretty :: a -> Text
     pretty = B.toLazyText . pretty'
+
+stdoutIsTTY :: Bool
+stdoutIsTTY = unsafePerformIO $ hIsTerminalDevice stdout
diff --git a/SourceFile.hs b/SourceFile.hs
index 0a054c0..8fc9dbf 100644
--- a/SourceFile.hs
+++ b/SourceFile.hs
@@ -19,19 +19,25 @@ data Occurrence = Occurrence
     , occArgs :: [(Text, Text)] }  -- prefix spacing, argument value
   deriving (Show)
 
-instance Pretty SourceFile where
-    pretty' (SourceFile cs) = mconcat (map pretty' cs)
-
-instance Pretty Chunk where
-    pretty' (CSkip s) = B.fromLazyText s
-    pretty' (COcc occ) = pretty' occ
-
-instance Pretty Occurrence where
-    pretty' (Occurrence name args) =
-        ansiBG '1' (B.fromLazyText name)
-        <> mconcat [ansiBG '4' (B.fromLazyText pre)
-                      <> ansiBG '3' (B.fromLazyText val)
-                   | (pre, val) <- args]
-      where
-        ansiBG :: Char -> Builder -> Builder
-        ansiBG c b = B.fromString ("\x1B[1;4" ++ c : "m") <> b <> B.fromString "\x1B[0m"
+occurrenceToBuilder :: Bool -> Occurrence -> Builder
+occurrenceToBuilder useClr (Occurrence name args) =
+    ansiBG '1' (B.fromLazyText name)
+    <> mconcat [ansiBG '4' (B.fromLazyText pre)
+                  <> ansiBG '3' (B.fromLazyText val)
+               | (pre, val) <- args]
+  where
+    ansiBG :: Char -> Builder -> Builder
+    ansiBG c b
+      | useClr && stdoutIsTTY = B.fromString ("\x1B[1;4" ++ c : "m") <> b <> B.fromString "\x1B[0m"
+      | otherwise = b
+
+chunkToBuilder :: Bool -> Chunk -> Builder
+chunkToBuilder _ (CSkip s) = B.fromLazyText s
+chunkToBuilder useClr (COcc occ) = occurrenceToBuilder useClr occ
+
+sourceFileToBuilder :: Bool -> SourceFile -> Builder
+sourceFileToBuilder useClr (SourceFile cs) = mconcat (map (chunkToBuilder useClr) cs)
+
+instance Pretty SourceFile where pretty' = sourceFileToBuilder True
+instance Pretty Chunk      where pretty' = chunkToBuilder      True
+instance Pretty Occurrence where pretty' = occurrenceToBuilder True
diff --git a/refactor-type-parameters.cabal b/refactor-type-parameters.cabal
index 63e7395..61b97c7 100644
--- a/refactor-type-parameters.cabal
+++ b/refactor-type-parameters.cabal
@@ -13,7 +13,8 @@ executable refactor-type-parameters
   build-depends:       base >= 4.13 && < 4.15,
                        mtl,
                        parsec,
-                       text
+                       text,
+                       optparse-applicative
   hs-source-dirs:      .
   default-language:    Haskell2010
   ghc-options:         -Wall -O2 -threaded -rtsopts