Remove shebang from input before attempting to extract extension pragmas.

2 files changed, 16 insertions, 6 deletions

diff --git a/lib/Language/Haskell/Stylish/Parse.hs b/lib/Language/Haskell/Stylish/Parse.hs
index f8e24a6..3118380 100644
--- a/lib/Language/Haskell/Stylish/Parse.hs
+++ b/lib/Language/Haskell/Stylish/Parse.hs
@@ -9,7 +9,6 @@ import           Data.Maybe                      (fromMaybe, listToMaybe)
 import qualified Language.Haskell.Exts.Annotated as H
 import Data.List (isPrefixOf)
 
-
 --------------------------------------------------------------------------------
 import           Language.Haskell.Stylish.Config
 import           Language.Haskell.Stylish.Step
@@ -48,9 +47,9 @@ dropBom str              = str
 parseModule :: Extensions -> Maybe FilePath -> String -> Either String Module
 parseModule extraExts mfp string = do
     -- Determine the extensions: those specified in the file and the extra ones
-    let noBom            = dropBom string
+    let noPrefixes       = unShebang . dropBom $ string
         extraExts'       = map H.classifyExtension extraExts
-        (lang, fileExts) = fromMaybe (Nothing, []) $ H.readExtensions noBom
+        (lang, fileExts) = fromMaybe (Nothing, []) $ H.readExtensions noPrefixes
         exts             = fileExts ++ extraExts'
 
         -- Parsing options...
@@ -64,8 +63,9 @@ parseModule extraExts mfp string = do
             }
 
         -- Preprocessing
-        processed = unShebang $
-            if H.EnableExtension H.CPP `elem` exts then unCpp noBom else noBom
+        processed = if H.EnableExtension H.CPP `elem` exts
+                       then unCpp noPrefixes
+                       else noPrefixes
 
     case H.parseModuleWithComments mode processed of
         H.ParseOk md -> return md
diff --git a/tests/Language/Haskell/Stylish/Parse/Tests.hs b/tests/Language/Haskell/Stylish/Parse/Tests.hs
index d5f3c1d..87c0a51 100644
--- a/tests/Language/Haskell/Stylish/Parse/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Parse/Tests.hs
@@ -1,4 +1,3 @@
---------------------------------------------------------------------------------
 module Language.Haskell.Stylish.Parse.Tests
     ( tests
     ) where
@@ -22,8 +21,19 @@ tests = testGroup "Language.Haskell.Stylish.Parse"
     , testCase "Multiline CPP"         testMultilineCpp
     , testCase "Haskell2010 extension" testHaskell2010
     , testCase "Shebang"               testShebang
+    , testCase "ShebangExt"            testShebangExt
     ]
 
+--------------------------------------------------------------------------------
+testShebangExt :: Assertion
+testShebangExt = assert $ isRight $ parseModule [] Nothing input
+                         where
+                           input = unlines
+                               [ "#!env runghc"
+                               , "{-# LANGUAGE CPP #-}"
+                               , "#define foo bar \\"
+                               , "             qux"
+                               ]
 
 --------------------------------------------------------------------------------
 testBom :: Assertion