Completely rewrite

nix-script.hs

@@ -1,131 +1,133 @@
-#!/usr/bin/env nix-script
+-- | A shebang for running scripts inside nix-shell with defined dependencies
-#!> haskell
+module NixScript where
-#! haskell | text lens
-#! shell   | nix
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TupleSections     #-}
 
-import Control.Monad
+import Control.Monad        (when)
-import Control.Applicative
+import Data.Maybe           (fromMaybe)
-import System.Environment
+import Data.List            (isSuffixOf, isPrefixOf, find)
-import Data.List
+import System.Environment   (lookupEnv, getProgName, getArgs)
-import Data.Text (Text)
+import System.Process       (callProcess)
-import Data.Text.Lens (_Text)
+import System.Posix.IO      (createPipe, fdToHandle)
-import Control.Lens
+import System.IO            (hPutStrLn, hClose, hFlush)
-import Control.Exception.Lens
-import System.IO.Error.Lens
-import System.Exit
-import System.Posix.Process
-import System.Posix.IO
-import System.IO
-import Data.Char
-import Data.Monoid
-
-import qualified Data.Text as Text
 
 
 -- | Information about a language
-data LangDef = LangDef
+data Language = Language
-  { name :: String                         -- ^ Name of this language
+  { name      :: String
-  , deps :: [Text]   -> [Text]             -- ^ Convert langunage-specific dependencies to nix packages
+    -- ^ Name of the language
-  , run  :: FilePath -> (String, [String]) -- ^ Command to run the given file as script
+  , depsTrans :: [String] -> [String]
-  , repl :: FilePath -> (String, [String]) -- ^ Command to load the given file in an interpreter
+    -- ^ Transform language-specific dependencies to nix packages
+  , run       :: FilePath -> (String, [String])
+    -- ^ Command to run the given file as script
+  , repl      :: FilePath -> (String, [String])
+    -- ^ Command to load the given file in an interpreter
   }
 
 
-basePackages :: [Text]
+-- | Basic packages always present
+basePackages :: [String]
 basePackages = ["coreutils", "utillinux"]
 
+-- | Preserved environment variables
+baseEnv :: [String]
+baseEnv = ["LOCALE_ARCHIVE", "LANG", "TERMINFO", "TERM"]
 
-languages :: [LangDef]
+
+-- | List of supported language definitions
+languages :: [Language]
 languages = [haskell, python, javascript, perl, shell]
-
+  where
-haskell :: LangDef
+    haskell = Language "haskell" d r i where
-haskell = LangDef "haskell" d r i where
+      d pkgs = pure ("haskellPackages.ghcWithPackages (hs: with hs; [" ++ 
-  d pkgs = pure $
+                     unwords pkgs ++ "])")
-    "haskellPackages.ghcWithPackages (hs: with hs; [" <> Text.unwords pkgs <> "])"
+      r script = ("runghc" , [script])
-  r script = ("runhaskell" , [script])
       i script = ("ghci"   , [script])
 
-python :: LangDef
+    python = Language "python" d r i where
-python = LangDef "python" d r i where
+      d pkgs   = "python" : map ("pythonPackages." ++) pkgs
-  d pkgs   = "python" : map ("pythonPackages." <>) pkgs
       r script = ("python" , [script])
       i script = ("python" , ["-i", script])
 
-javascript :: LangDef
+    javascript = Language "javascript" d r i where
-javascript = LangDef "javascript" d r i where
+      d pkgs   = "node" : map ("nodePackages." ++) pkgs
-  d pkgs   = "node" : map ("nodePackages." <>) pkgs
       r script = ("node" , [script])
       i script = ("node" , [])
 
-perl :: LangDef
+    perl = Language "perl" d r i where
-perl = LangDef "perl" d r i where
+      d pkgs   = "perl" : map ("perlPackages." ++) pkgs
-  d pkgs   = "perl" : map ("perlPackages." <>) pkgs
       r script = ("perl" , [script])
       i script = ("perl" , ["-d", script])
 
-shell :: LangDef
+    shell = Language "shell" d r i where
-shell = LangDef "shell" d r i where
       d        = mappend ("bash" : basePackages)
       r script = ("bash", [script])
       i _      = ("bash", [])
 
-passthrough :: String -> LangDef
+
-passthrough name = LangDef name d r i where
+-- | Create ad-hoc definitions for unknown languages
+passthrough :: String -> Language
+passthrough name = Language name d r i where
   d        = mappend basePackages
   r script = (name, [script])
   i _      = (name, [])
 
 
-lookupLangDef :: String -> IO LangDef
+-- | Find the appropriate language definition
-lookupLangDef n
+lookupLang :: String -> Language
-  | Just def <- find ((n ==) . name) languages = return def
+lookupLang n =
-  | otherwise = return (passthrough n)
+  fromMaybe (passthrough n) (find ((n ==) . name) languages)
 
-makeDeps :: String -> [String] -> IO [String]
+-- | Parse dependencies declaration line
-makeDeps lang ds = lookupLangDef lang <&> \def ->
+parseHeader :: String -> [String]
-  map (view _Text) $ deps def (map (review _Text) ds)
+parseHeader = uncurry trans . split . words
+  where
+    trans lang = depsTrans (lookupLang lang)
+    split (lang : "|" : deps) = (lang, deps)
+    split line = error ("Invalid dependency declaration: " ++ unwords line)
 
-parseDepLine :: [String] -> IO (String, [String])
-parseDepLine (lang:"|":deps) = return (lang, deps)
-parseDepLine x = fail $ "Invalid dependency specification: " ++ unwords x
 
-makeCommand :: String -> Bool -> String -> IO (String, [String])
+-- | Find command to run/load the script
-makeCommand lang interactive file = lookupLangDef lang <&> \def ->
+interpreter :: String -> Bool -> String -> (String, [String])
-  (if interactive then repl else run) def file
+interpreter lang interactive = 
+  (if interactive then repl else run) (lookupLang lang)
 
-makeEnvArg :: String -> IO String
-makeEnvArg env = f $ getEnv env <&> \val -> env ++ "=" ++ val where
-  f = handling_ (_IOException.errorType._NoSuchThing) $ return ""
 
-makeXargsCommand :: String -> Int -> IO String
+-- | Create command to add the shell environment
-makeXargsCommand cmd fd = do
+makeCommand :: String -> [String] -> IO String
-  let xargsFile = "/proc/self/fd/" ++ show fd
+makeCommand program args = do
-  envStr <- unwords <$> traverse makeEnvArg
+  (readFd, writeFd) <- createPipe
-    ["LOCALE_ARCHIVE", "LANG", "TERMINFO", "TERM"]
+  writeH <- fdToHandle writeFd
-  return $ "env " ++ envStr ++ " xargs -a " ++ xargsFile ++ " -d '\\n' " ++ cmd ++ ""
+  hPutStrLn writeH (unlines args)
+  hFlush writeH >> hClose writeH
 
+  definitions <- mapM format baseEnv
+  return (env definitions ++ xargs readFd ++ program)
+  where
+    env defs   = "env " ++ unwords defs ++ " "
+    xargs fd   = "xargs -a /proc/self/fd/" ++ show fd ++ " -d '\\n' "
+    format var = maybe "" (\x -> var ++ "=" ++ x) <$> lookupEnv var
+
+
+-- | run a script or load it in an interactive interpreter
 main :: IO ()
 main = do
   progName <- getProgName
-  args <- getArgs
+  progArgs <- getArgs
-  let interactive = "i" `isSuffixOf` progName
+
-  case args ^? _Cons of
+  when (null progArgs) (fail $ "usage: " ++ progName ++ " <file>")
-    Nothing -> fail $ "usage: " ++ progName ++ " <file>" ++ "   [missing file name]"
+
-    Just (file, args') -> do
+  let file    = head progArgs
-      header <- drop 1 . map (drop 2) . takeWhile ("#!" `isPrefixOf`) . lines <$> readFile file
+      shebang = takeWhile (isPrefixOf "#!") . lines
-      case header ^? _Cons of
+      header  = drop 1  . map (drop 2) . shebang
-        Just ('>':lang, depHeader) -> do
+
-          deps   <- concat <$> traverse (uncurry makeDeps <=< parseDepLine . words) depHeader
+  script <- readFile file
-          let deps' = "findutils" : deps
+  case header script of
-          let depArgs = concatMap (\x -> ["-p", x]) deps'
+    (('>' : identifier) : lines) -> do
-          (cmd,cmdArgs) <- makeCommand (under _Text Text.strip lang) interactive file
+      let pkgs            = concatMap parseHeader lines
-          (readFd, writeFd) <- createPipe
+          language        = dropWhile (==' ') identifier
-          writeH <- fdToHandle writeFd
+          interactive     = isSuffixOf "i" progName
-          hPutStrLn writeH (unlines cmdArgs) >> hFlush writeH
+          (program, args) = interpreter language interactive file
-          hClose writeH
+
-          xargsCmd <- makeXargsCommand cmd (fromIntegral readFd)
+      cmd <- makeCommand program args
-          let finalArgs = "--pure" : "--command" : xargsCmd : depArgs
+      callProcess "nix-shell" ("--pure" : "--command" : cmd : "-p" : pkgs)
-          executeFile "nix-shell" True finalArgs Nothing
+
-        _ -> fail "missing language to run as"
+    _ -> fail "missing or invalid header"