Merge pull request #351 from github/semantic-ast-errors

Semantic AST polish

semantic-ast/src/Main.hs

@@ -3,8 +3,8 @@ module Main (main) where
 
 import System.Environment
 import TreeSitter.Unmarshal
-import TreeSitter.Python.AST
-import TreeSitter.Python
+import qualified TreeSitter.Python.AST as AST
+import qualified TreeSitter.Python as Python
 import Source.Range
 import Source.Span
 import Data.ByteString.Char8
@@ -12,42 +12,53 @@ import Data.ByteString (pack, readFile, ByteString)
 import System.IO (FilePath)
 import Options.Applicative hiding (style)
 import Data.Semigroup ((<>))
-import Text.Pretty.Simple (pPrint)
+import Text.Pretty.Simple (pPrint, pPrintNoColor)
+import Data.Foldable (traverse_)
+import Control.Monad ((>=>))
 
 data SemanticAST = SemanticAST
   { format :: Format
-  , source :: Either FilePath Prelude.String
+  , noColor  :: Bool
+  , source :: Either [FilePath] String
   }
 
+-- Usage: semantic-ast --format ARG [--no-color] (--sourceString STRING | FILEPATHS…)
 parseAST :: Parser SemanticAST
 parseAST = SemanticAST
   <$> option auto
     ( long "format"
     <> help "Specify desired output: show, json, sexpression" )
-  <*> (Left <$> strOption
-      ( long "sourceFile"
-      <> metavar "FILEPATH"
-      <> help "Specify filepath containing source code to parse" )
+  <*> switch
+    ( long "no-color"
+    <> help "Print with color: --color"
+    )
+  <*> (Left <$> some
+    (Options.Applicative.argument str (metavar "FILEPATH(S)"))
     <|> Right <$> strOption
       ( long "sourceString"
       <> metavar "STRING"
       <> help "Specify source input to parse"
       ))
 
+
 main :: IO ()
 main = generateAST =<< execParser opts
 
+
 generateAST :: SemanticAST -> IO ()
-generateAST (SemanticAST format source) = do
-  bytestring <- case source of
-    Left filePath -> do
-      Data.ByteString.readFile filePath
-    Right source -> do
-      pure $ Data.ByteString.Char8.pack source
-  ast <- parseByteString @TreeSitter.Python.AST.Module @(Range, Span) tree_sitter_python bytestring
-  case format of
-    Show -> print ast
-    Pretty -> pPrint ast
+generateAST (SemanticAST format noColor source) =
+  getByteStrings >>= traverse_ go
+  where getByteStrings = case source of
+          Left filePaths -> traverse Data.ByteString.readFile filePaths
+          Right source   -> pure [Data.ByteString.Char8.pack source]
+        go = ast >=> display
+        ast = parseByteString @AST.Module @(Range, Span) Python.tree_sitter_python
+        display = case format of
+          Show -> print
+          Pretty | noColor -> pPrintNoColor
+                 | otherwise -> pPrint
+
+-- need AST in scope for case format and ..
 
 opts :: ParserInfo SemanticAST
 opts = info (parseAST <**> helper)
@@ -56,5 +67,8 @@ opts = info (parseAST <**> helper)
  <> header "semantic-ast is a package used to parse source code" )
 
 -- TODO: Define formats for json, sexpression, etc.
-data Format = Show | Pretty
+data Format = Show
+            | Pretty
   deriving (Read)
+
+-- bool field would break Read