Merge pull request #143 from themoritz/fix/keep-safe-and-source

Keep `safe` and `{-# SOURCE #-}` import annotations

2 files changed, 35 insertions, 1 deletions

diff --git a/lib/Language/Haskell/Stylish/Step/Imports.hs b/lib/Language/Haskell/Stylish/Step/Imports.hs
index 334fa15..29b8cc2 100644
--- a/lib/Language/Haskell/Stylish/Step/Imports.hs
+++ b/lib/Language/Haskell/Stylish/Step/Imports.hs
@@ -235,6 +235,8 @@ prettyImport columns Options{..} padQualified padName longest imp
 
     base' baseName importAs hasHiding' = unwords $ concat $ filter (not . null)
         [ ["import"]
+        , source
+        , safe
         , qualified
         , show <$> maybeToList (H.importPkg imp)
         , [baseName]
@@ -259,9 +261,22 @@ prettyImport columns Options{..} padQualified padName longest imp
 
     qualified
         | H.importQualified imp = ["qualified"]
-        | padQualified          = ["         "]
+        | padQualified          =
+              if H.importSrc imp
+                  then []
+                  else if H.importSafe imp
+                           then ["    "]
+                           else ["         "]
         | otherwise             = []
 
+    safe
+        | H.importSafe imp = ["safe"]
+        | otherwise        = []
+
+    source
+        | H.importSrc imp = ["{-# SOURCE #-}"]
+        | otherwise       = []
+
     mapSpecs f = case importSpecs of
         Nothing -> []     -- Import everything
         Just [] -> ["()"] -- Instance only imports
diff --git a/tests/Language/Haskell/Stylish/Step/Imports/Tests.hs b/tests/Language/Haskell/Stylish/Step/Imports/Tests.hs
index 5ca60de..c3178ac 100644
--- a/tests/Language/Haskell/Stylish/Step/Imports/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Step/Imports/Tests.hs
@@ -46,6 +46,7 @@ tests = testGroup "Language.Haskell.Stylish.Step.Imports.Tests"
     , testCase "case 19b" case19b
     , testCase "case 19d" case19c
     , testCase "case 19d" case19d
+    , testCase "case 20" case20
     ]
 
 
@@ -512,3 +513,21 @@ case19input = unlines
         , ""
         , "import Data.List (foldl', intercalate, intersperse)"
         ]
+
+--------------------------------------------------------------------------------
+case20 :: Assertion
+case20 = expected
+    @=? testStep (step 80 defaultOptions) input'
+  where
+    expected = unlines
+        [ "import {-# SOURCE #-} Data.ByteString as BS"
+        , "import qualified Data.Map        as Map"
+        , "import           Data.Set        (empty)"
+        , "import {-# SOURCE #-} qualified Data.Text       as T"
+        ]
+    input' = unlines
+        [ "import {-# SOURCE #-}    Data.ByteString as BS"
+        , "import {-# SOURCE #-} qualified Data.Text as T"
+        , "import qualified   Data.Map as Map"
+        , "import Data.Set (empty)"
+        ]