Align option for LanguagePragmas

5 files changed, 87 insertions, 26 deletions

diff --git a/data/stylish-haskell.yaml b/data/stylish-haskell.yaml
index 0778f93..3d4e98a 100644
--- a/data/stylish-haskell.yaml
+++ b/data/stylish-haskell.yaml
@@ -94,6 +94,16 @@ steps:
       # Default: vertical.
       style: vertical
 
+      # Align affects alignment of closing pragma brackets.
+      #
+      # - true: Brackets are aligned in same collumn.
+      #
+      # - false: Brackets are not aligned together. There is only one space
+      #   between actual import and closing bracket.
+      #
+      # Default: true
+      align: true
+
       # stylish-haskell can detect redundancy of some language pragmas. If this
       # is set to true, it will remove those redundant pragmas. Default: true.
       remove_redundant: true
diff --git a/src/Language/Haskell/Stylish.hs b/src/Language/Haskell/Stylish.hs
index 7a52aa2..b8620ae 100644
--- a/src/Language/Haskell/Stylish.hs
+++ b/src/Language/Haskell/Stylish.hs
@@ -53,6 +53,7 @@ imports = Imports.step
 --------------------------------------------------------------------------------
 languagePragmas :: Int -- ^ columns
                 -> LanguagePragmas.Style
+                -> Bool -- ^ Pad to same length in vertical mode?
                 -> Bool -- ^ remove redundant?
                 -> Step
 languagePragmas = LanguagePragmas.step
diff --git a/src/Language/Haskell/Stylish/Config.hs b/src/Language/Haskell/Stylish/Config.hs
index 72b59ca..c176eb6 100644
--- a/src/Language/Haskell/Stylish/Config.hs
+++ b/src/Language/Haskell/Stylish/Config.hs
@@ -197,6 +197,7 @@ parseLanguagePragmas :: Config -> A.Object -> A.Parser Step
 parseLanguagePragmas config o = LanguagePragmas.step
     <$> pure (configColumns config)
     <*> (o A..:? "style" >>= parseEnum styles LanguagePragmas.Vertical)
+    <*> o A..:? "align" A..!= True
     <*> o A..:? "remove_redundant" A..!= True
   where
     styles =
diff --git a/src/Language/Haskell/Stylish/Step/LanguagePragmas.hs b/src/Language/Haskell/Stylish/Step/LanguagePragmas.hs
index 209b2f2..0239736 100644
--- a/src/Language/Haskell/Stylish/Step/LanguagePragmas.hs
+++ b/src/Language/Haskell/Stylish/Step/LanguagePragmas.hs
@@ -42,11 +42,15 @@ firstLocation = minimum . map (blockStart . fst)
 
 
 --------------------------------------------------------------------------------
-verticalPragmas :: Int -> [String] -> Lines
-verticalPragmas longest pragmas' =
-    [ "{-# LANGUAGE " ++ padRight longest pragma ++ " #-}"
+verticalPragmas :: Int -> Bool -> [String] -> Lines
+verticalPragmas longest align pragmas' =
+    [ "{-# LANGUAGE " ++ pad pragma ++ " #-}"
     | pragma <- pragmas'
     ]
+  where
+    pad
+      | align = padRight longest
+      | otherwise = id
 
 
 --------------------------------------------------------------------------------
@@ -56,17 +60,23 @@ compactPragmas columns pragmas' = wrap columns "{-# LANGUAGE" 13 $
 
 
 --------------------------------------------------------------------------------
-compactLinePragmas :: Int -> [String] -> Lines
-compactLinePragmas _ [] = []
-compactLinePragmas columns pragmas' =
-  let maxWidth = columns - 16
-      longest  = maximum $ map length prags
-      prags    = map truncateComma $ wrap maxWidth "" 1 $
-                    map (++ ",") (init pragmas') ++ [last pragmas']
-  in map (wrapLanguage . padRight longest) prags
+compactLinePragmas :: Int -> Bool -> [String] -> Lines
+compactLinePragmas _ _ [] = []
+compactLinePragmas columns align pragmas' = map (wrapLanguage . pad) prags
   where
     wrapLanguage ps = "{-# LANGUAGE" ++ ps ++  " #-}"
 
+    maxWidth = columns - 16
+
+    longest  = maximum $ map length prags
+
+    pad
+      | align = padRight longest
+      | otherwise = id
+
+    prags = map truncateComma $ wrap maxWidth "" 1 $
+      map (++ ",") (init pragmas') ++ [last pragmas']
+
 
 --------------------------------------------------------------------------------
 truncateComma :: String -> String
@@ -77,10 +87,10 @@ truncateComma xs
 
 
 --------------------------------------------------------------------------------
-prettyPragmas :: Int -> Int -> Style -> [String] -> Lines
-prettyPragmas _       longest Vertical    = verticalPragmas longest
-prettyPragmas columns _       Compact     = compactPragmas columns
-prettyPragmas columns _       CompactLine = compactLinePragmas columns
+prettyPragmas :: Int -> Int -> Bool -> Style -> [String] -> Lines
+prettyPragmas _    longest align Vertical    = verticalPragmas longest align
+prettyPragmas cols _       _     Compact     = compactPragmas cols
+prettyPragmas cols _       align CompactLine = compactLinePragmas cols align
 
 
 --------------------------------------------------------------------------------
@@ -100,13 +110,13 @@ filterRedundant isRedundant' = snd . foldr filterRedundant' (S.empty, [])
         known' = xs' `S.union` known
 
 --------------------------------------------------------------------------------
-step :: Int -> Style -> Bool -> Step
-step columns style = makeStep "LanguagePragmas" . step' columns style
+step :: Int -> Style -> Bool -> Bool -> Step
+step = (((makeStep "LanguagePragmas" .) .) .) . step'
 
 
 --------------------------------------------------------------------------------
-step' :: Int -> Style -> Bool -> Lines -> Module -> Lines
-step' columns style removeRedundant ls (module', _)
+step' :: Int -> Style -> Bool -> Bool -> Lines -> Module -> Lines
+step' columns style align removeRedundant ls (module', _)
     | null pragmas' = ls
     | otherwise     = applyChanges changes ls
   where
@@ -118,7 +128,7 @@ step' columns style removeRedundant ls (module', _)
     longest  = maximum $ map length $ snd =<< pragmas'
     groups   = [(b, concat pgs) | (b, pgs) <- groupAdjacent pragmas']
     changes  =
-        [ change b (const $ prettyPragmas columns longest style pg)
+        [ change b (const $ prettyPragmas columns longest align style pg)
         | (b, pg) <- filterRedundant isRedundant' groups
         ]
 
diff --git a/tests/Language/Haskell/Stylish/Step/LanguagePragmas/Tests.hs b/tests/Language/Haskell/Stylish/Step/LanguagePragmas/Tests.hs
index 3cfabef..fe889e4 100644
--- a/tests/Language/Haskell/Stylish/Step/LanguagePragmas/Tests.hs
+++ b/tests/Language/Haskell/Stylish/Step/LanguagePragmas/Tests.hs
@@ -24,12 +24,14 @@ tests = testGroup "Language.Haskell.Stylish.Step.LanguagePragmas.Tests"
     , testCase "case 04" case04
     , testCase "case 05" case05
     , testCase "case 06" case06
+    , testCase "case 07" case07
+    , testCase "case 08" case08
     ]
 
 
 --------------------------------------------------------------------------------
 case01 :: Assertion
-case01 = expected @=? testStep (step 80 Vertical False) input
+case01 = expected @=? testStep (step 80 Vertical True False) input
   where
     input = unlines
         [ "{-# LANGUAGE ViewPatterns #-}"
@@ -48,7 +50,7 @@ case01 = expected @=? testStep (step 80 Vertical False) input
 
 --------------------------------------------------------------------------------
 case02 :: Assertion
-case02 = expected @=? testStep (step 80 Vertical True) input
+case02 = expected @=? testStep (step 80 Vertical True True) input
   where
     input = unlines
         [ "{-# LANGUAGE BangPatterns #-}"
@@ -64,7 +66,7 @@ case02 = expected @=? testStep (step 80 Vertical True) input
 
 --------------------------------------------------------------------------------
 case03 :: Assertion
-case03 = expected @=? testStep (step 80 Vertical True) input
+case03 = expected @=? testStep (step 80 Vertical True True) input
   where
     input = unlines
         [ "{-# LANGUAGE BangPatterns #-}"
@@ -80,7 +82,7 @@ case03 = expected @=? testStep (step 80 Vertical True) input
 
 --------------------------------------------------------------------------------
 case04 :: Assertion
-case04 = expected @=? testStep (step 80 Compact False) input
+case04 = expected @=? testStep (step 80 Compact True False) input
   where
     input = unlines
         [ "{-# LANGUAGE TypeOperators, StandaloneDeriving, DeriveDataTypeable,"
@@ -97,7 +99,7 @@ case04 = expected @=? testStep (step 80 Compact False) input
 
 --------------------------------------------------------------------------------
 case05 :: Assertion
-case05 = expected @=? testStep (step 80 Vertical False) input
+case05 = expected @=? testStep (step 80 Vertical True False) input
   where
     input = unlines
         [ "{-# LANGUAGE CPP #-}"
@@ -115,8 +117,10 @@ case05 = expected @=? testStep (step 80 Vertical False) input
         , "#endif"
         ]
 
+
+--------------------------------------------------------------------------------
 case06 :: Assertion
-case06 = expected @=? testStep (step 80 CompactLine False) input
+case06 = expected @=? testStep (step 80 CompactLine True False) input
   where
     input = unlines
         [ "{-# LANGUAGE TypeOperators, StandaloneDeriving, DeriveDataTypeable,"
@@ -128,3 +132,38 @@ case06 = expected @=? testStep (step 80 CompactLine False) input
           "TemplateHaskell #-}"
         , "{-# LANGUAGE TypeOperators, ViewPatterns                             #-}"
         ]
+
+--------------------------------------------------------------------------------
+case07 :: Assertion
+case07 = expected @=? testStep (step 80 Vertical False False) input
+  where
+    input = unlines
+        [ "{-# LANGUAGE ViewPatterns #-}"
+        , "{-# LANGUAGE TemplateHaskell, ViewPatterns #-}"
+        , "{-# LANGUAGE ScopedTypeVariables, NoImplicitPrelude #-}"
+        , "module Main where"
+        ]
+
+    expected = unlines
+        [ "{-# LANGUAGE NoImplicitPrelude #-}"
+        , "{-# LANGUAGE ScopedTypeVariables #-}"
+        , "{-# LANGUAGE TemplateHaskell #-}"
+        , "{-# LANGUAGE ViewPatterns #-}"
+        , "module Main where"
+        ]
+
+
+--------------------------------------------------------------------------------
+case08 :: Assertion
+case08 = expected @=? testStep (step 80 CompactLine False False) input
+  where
+    input = unlines
+        [ "{-# LANGUAGE TypeOperators, StandaloneDeriving, DeriveDataTypeable,"
+        , "    TemplateHaskell #-}"
+        , "{-# LANGUAGE TemplateHaskell, ViewPatterns #-}"
+        ]
+    expected = unlines
+        [ "{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving, " ++
+          "TemplateHaskell #-}"
+        , "{-# LANGUAGE TypeOperators, ViewPatterns #-}"
+        ]