Align record data declarations

6 files changed, 81 insertions, 7 deletions

diff --git a/.stylish-haskell.yaml b/.stylish-haskell.yaml
index 4be6dc7..c128830 100644
--- a/.stylish-haskell.yaml
+++ b/.stylish-haskell.yaml
@@ -45,6 +45,9 @@ steps:
       # is set to true, it will remove those redundant pragmas. Default: true.
       remove_redundant: true
 
+  # Align the types in record declarations
+  - records: {}
+
   # Replace tabs by spaces. This is disabled by default.
   # - tabs:
   #     # Number of spaces to use for each tab. Default: 8, as specified by the
diff --git a/src/StylishHaskell/Config.hs b/src/StylishHaskell/Config.hs
index e2e6869..9d39ccd 100644
--- a/src/StylishHaskell/Config.hs
+++ b/src/StylishHaskell/Config.hs
@@ -12,7 +12,7 @@ module StylishHaskell.Config
 --------------------------------------------------------------------------------
 import           Control.Applicative                    (pure, (<$>), (<*>))
 import           Control.Monad                          (forM, msum, mzero)
-import           Data.Aeson                             (FromJSON(..))
+import           Data.Aeson                             (FromJSON (..))
 import qualified Data.Aeson                             as A
 import qualified Data.Aeson.Types                       as A
 import qualified Data.ByteString                        as B
@@ -29,6 +29,7 @@ import           Paths_stylish_haskell                  (getDataFileName)
 import           StylishHaskell.Step
 import qualified StylishHaskell.Step.Imports            as Imports
 import qualified StylishHaskell.Step.LanguagePragmas    as LanguagePragmas
+import qualified StylishHaskell.Step.Records            as Records
 import qualified StylishHaskell.Step.Tabs               as Tabs
 import qualified StylishHaskell.Step.TrailingWhitespace as TrailingWhitespace
 import qualified StylishHaskell.Step.UnicodeSyntax      as UnicodeSyntax
@@ -124,6 +125,7 @@ catalog :: Map String (Config -> A.Object -> A.Parser Step)
 catalog = M.fromList
     [ ("imports",             parseImports)
     , ("language_pragmas",    parseLanguagePragmas)
+    , ("records",             parseRecords)
     , ("tabs",                parseTabs)
     , ("trailing_whitespace", parseTrailingWhitespace)
     , ("unicode_syntax",      parseUnicodeSyntax)
@@ -176,6 +178,11 @@ parseLanguagePragmas config o = LanguagePragmas.step
 
 
 --------------------------------------------------------------------------------
+parseRecords :: Config -> A.Object -> A.Parser Step
+parseRecords _ _ = return Records.step
+
+
+--------------------------------------------------------------------------------
 parseTabs :: Config -> A.Object -> A.Parser Step
 parseTabs _ o = Tabs.step
     <$> o A..:? "spaces" A..!= 8
diff --git a/src/StylishHaskell/Step/Records.hs b/src/StylishHaskell/Step/Records.hs
index ae9e0c4..f9c7348 100644
--- a/src/StylishHaskell/Step/Records.hs
+++ b/src/StylishHaskell/Step/Records.hs
@@ -1,15 +1,18 @@
 --------------------------------------------------------------------------------
-module StylishHaskell.Step.Records where
+module StylishHaskell.Step.Records
+    ( step
+    ) where
 
 
 --------------------------------------------------------------------------------
-import           Control.Arrow                   (second)
+import           Data.Char                       (isSpace)
 import           Data.List                       (nub)
 import qualified Language.Haskell.Exts.Annotated as H
 
 
 --------------------------------------------------------------------------------
 import           StylishHaskell.Editor
+import           StylishHaskell.Step
 import           StylishHaskell.Util
 
 
@@ -24,27 +27,35 @@ records modu =
 
 
 --------------------------------------------------------------------------------
--- | Align the types of a field
+-- | Align the type of a field
 alignType :: Int -> H.FieldDecl H.SrcSpan -> Change String
 alignType longest (H.FieldDecl srcSpan _ _) =
     changeLine (H.srcSpanStartLine srcSpan) alignType'
   where
     alignType' str =
-        let (pre, post) = second (drop 2) $ break (== ':') str
-        in [padRight longest pre ++ post]
+        let (pre, post) = break (== ':') str
+        in [padRight longest (trimRight pre) ++ post]
+
+    trimRight = reverse . dropWhile isSpace . reverse
 
 
 --------------------------------------------------------------------------------
 -- | Find the length of the longest field name in a record
 longestFieldName :: [H.FieldDecl H.SrcSpan] -> Int
 longestFieldName fields = maximum
-    [ H.srcSpanEndColumn $ H.ann name
+    [ H.srcSpanEndColumn (H.ann name)
     | H.FieldDecl _ names _ <- fields
     , name                  <- names
     ]
 
 
 --------------------------------------------------------------------------------
+-- | Align all fields in a record
+alignRecord :: [H.FieldDecl H.SrcSpan] -> [Change String]
+alignRecord fields = map (alignType $ longestFieldName fields) fields
+
+
+--------------------------------------------------------------------------------
 -- | Checks that all no field of the record appears on more than one line,
 -- amonst other things
 fixable :: [H.FieldDecl H.SrcSpan] -> Bool
@@ -54,3 +65,11 @@ fixable fields = all singleLine srcSpans && nonOverlapping srcSpans
     srcSpans          = map H.ann fields
     singleLine s      = H.srcSpanStartLine s == H.srcSpanEndLine s
     nonOverlapping ss = length ss == length (nub $ map H.srcSpanStartLine ss)
+
+
+--------------------------------------------------------------------------------
+step :: Step
+step = makeStep "Records" $ \ls (module', _) ->
+    let module''       = fmap H.srcInfoSpan module'
+        fixableRecords = filter fixable $ records module''
+    in applyChanges (fixableRecords >>= alignRecord) ls
diff --git a/stylish-haskell.cabal b/stylish-haskell.cabal
index 4aec024..684d798 100644
--- a/stylish-haskell.cabal
+++ b/stylish-haskell.cabal
@@ -36,6 +36,7 @@ Executable stylish-haskell
     StylishHaskell.Step
     StylishHaskell.Step.Imports
     StylishHaskell.Step.LanguagePragmas
+    StylishHaskell.Step.Records
     StylishHaskell.Step.Tabs
     StylishHaskell.Step.TrailingWhitespace
     StylishHaskell.Step.UnicodeSyntax
@@ -65,6 +66,7 @@ Test-suite stylish-haskell-tests
   Other-modules:
     StylishHaskell.Step.Imports.Tests
     StylishHaskell.Step.LanguagePragmas.Tests
+    StylishHaskell.Step.Records.Tests
     StylishHaskell.Step.Tabs.Tests
     StylishHaskell.Step.TrailingWhitespace.Tests
     StylishHaskell.Step.UnicodeSyntax.Tests
diff --git a/tests/StylishHaskell/Step/Records/Tests.hs b/tests/StylishHaskell/Step/Records/Tests.hs
new file mode 100644
index 0000000..4a1f92d
--- /dev/null
+++ b/tests/StylishHaskell/Step/Records/Tests.hs
@@ -0,0 +1,41 @@
+--------------------------------------------------------------------------------
+module StylishHaskell.Step.Records.Tests
+    ( tests
+    ) where
+
+
+--------------------------------------------------------------------------------
+import           Test.Framework                 (Test, testGroup)
+import           Test.Framework.Providers.HUnit (testCase)
+import           Test.HUnit                     (Assertion, (@=?))
+
+
+--------------------------------------------------------------------------------
+import           StylishHaskell.Step.Records
+import           StylishHaskell.Tests.Util
+
+
+--------------------------------------------------------------------------------
+tests :: Test
+tests = testGroup "StylishHaskell.Step.Records.Tests"
+    [ testCase "case 01" case01
+    ]
+
+
+--------------------------------------------------------------------------------
+case01 :: Assertion
+case01 = expected @=? testStep step input
+  where
+    input = unlines
+        [ "data Foo = Foo"
+        , "    { foo :: Int"
+        , "    , barqux :: String"
+        , "    } deriving (Show)"
+        ]
+
+    expected = unlines
+        [ "data Foo = Foo"
+        , "    { foo    :: Int"
+        , "    , barqux :: String"
+        , "    } deriving (Show)"
+        ]
diff --git a/tests/TestSuite.hs b/tests/TestSuite.hs
index e1c8acb..f5a7bff 100644
--- a/tests/TestSuite.hs
+++ b/tests/TestSuite.hs
@@ -12,6 +12,7 @@ import           Test.Framework                               (defaultMain)
 import qualified StylishHaskell.Parse.Tests
 import qualified StylishHaskell.Step.Imports.Tests
 import qualified StylishHaskell.Step.LanguagePragmas.Tests
+import qualified StylishHaskell.Step.Records.Tests
 import qualified StylishHaskell.Step.Tabs.Tests
 import qualified StylishHaskell.Step.TrailingWhitespace.Tests
 import qualified StylishHaskell.Step.UnicodeSyntax.Tests
@@ -23,6 +24,7 @@ main = defaultMain
     [ StylishHaskell.Parse.Tests.tests
     , StylishHaskell.Step.Imports.Tests.tests
     , StylishHaskell.Step.LanguagePragmas.Tests.tests
+    , StylishHaskell.Step.Records.Tests.tests
     , StylishHaskell.Step.Tabs.Tests.tests
     , StylishHaskell.Step.TrailingWhitespace.Tests.tests
     , StylishHaskell.Step.UnicodeSyntax.Tests.tests