Wrap import specs

3 files changed, 55 insertions, 17 deletions

diff --git a/src/StylishHaskell/Step/Imports.hs b/src/StylishHaskell/Step/Imports.hs
index 730006b..25be81e 100644
--- a/src/StylishHaskell/Step/Imports.hs
+++ b/src/StylishHaskell/Step/Imports.hs
@@ -6,6 +6,7 @@ module StylishHaskell.Step.Imports
 
 
 --------------------------------------------------------------------------------
+import           Control.Applicative             ((<$>))
 import           Control.Arrow                   ((&&&))
 import           Data.Char                       (isAlpha, toLower)
 import           Data.List                       (sortBy)
@@ -93,19 +94,34 @@ sortImportSpecs imp = imp {H.importSpecs = fmap sort $ H.importSpecs imp}
 --------------------------------------------------------------------------------
 prettyImport :: Bool -> Bool -> Int -> H.ImportDecl l -> String
 prettyImport padQualified padName longest imp =
-  unwords $ base : maybeToList specs
+    unlines specs
   where
     base = unwords $ concat
          [ ["import"]
          , qualified
          , [(if hasExtras && padName then padRight longest else id)
             (importName imp)]
-         , ["as " ++ as | H.ModuleName _ as <- maybeToList $ H.importAs imp]]
-    specs = fmap (unlines . padSpecs . lines  . H.prettyPrint)
-            $ H.importSpecs imp
-    pad = replicate (length base) ' '
-    padSpecs [] = []
-    padSpecs (h:rest) = h : map (pad ++) rest
+         , ["as " ++ as | H.ModuleName _ as <- maybeToList $ H.importAs imp]
+         ]
+
+
+    -- specs = unlines . indentSpecs . lines . H.prettyPrint <$> H.importSpecs imp
+    specs = wrap 80 base (length base + 1) $
+        (if hiding then ("hiding" :) else id) $
+        withInit (++ ",") $
+        withHead ("(" ++) $
+        withLast (++ ")") $
+        map H.prettyPrint $
+        importSpecs
+
+    hiding = case H.importSpecs imp of
+        Just (H.ImportSpecList _ h _) -> h
+        _                             -> False
+
+    importSpecs = case H.importSpecs imp of
+        Just (H.ImportSpecList _ _ l) -> l
+        _                             -> []
+
     hasExtras = isJust (H.importAs imp) || isJust (H.importSpecs imp)
 
     qualified
@@ -114,6 +130,7 @@ prettyImport padQualified padName longest imp =
         | otherwise             = []
 
 
+
 --------------------------------------------------------------------------------
 prettyImportGroup :: Align -> Int -> [H.ImportDecl LineBlock] -> Lines
 prettyImportGroup align longest imps =
diff --git a/src/StylishHaskell/Step/UnicodeSyntax.hs b/src/StylishHaskell/Step/UnicodeSyntax.hs
index e833107..3311b0e 100644
--- a/src/StylishHaskell/Step/UnicodeSyntax.hs
+++ b/src/StylishHaskell/Step/UnicodeSyntax.hs
@@ -103,13 +103,6 @@ between (startRow, startCol) (endRow, endCol) needle =
     take (endRow - startRow + 1) .
     drop (startRow - 1)
   where
-    withHead _ []       = []
-    withHead f (x : xs) = (f x) : xs
-
-    withLast f [x]      = [f x]
-    withLast f (x : xs) = x : withLast f xs
-    withLast _ []       = []
-
     search _      []            = Nothing
     search (r, _) ([] : xs)     = search (r + 1, 1) xs
     search (r, c) (x : xs)
diff --git a/src/StylishHaskell/Util.hs b/src/StylishHaskell/Util.hs
index 52e6323..2dde009 100644
--- a/src/StylishHaskell/Util.hs
+++ b/src/StylishHaskell/Util.hs
@@ -1,10 +1,15 @@
 --------------------------------------------------------------------------------
 module StylishHaskell.Util
     ( nameToString
+    , indent
     , padRight
     , everything
     , infoPoints
     , wrap
+
+    , withHead
+    , withLast
+    , withInit
     ) where
 
 
@@ -28,6 +33,11 @@ nameToString (H.Symbol _ str) = str
 
 
 --------------------------------------------------------------------------------
+indent :: Int -> String -> String
+indent len str = replicate len ' ' ++ str
+
+
+--------------------------------------------------------------------------------
 padRight :: Int -> String -> String
 padRight len str = str ++ replicate (len - length str) ' '
 
@@ -48,17 +58,35 @@ wrap :: Int       -- ^ Maximum line width
      -> Int       -- ^ Indentation
      -> [String]  -- ^ Strings to add/wrap
      -> Lines     -- ^ Resulting lines
-wrap maxWidth leading indent strs =
+wrap maxWidth leading ind strs =
     let (ls, curr, _) = foldl step ([], leading, length leading) strs
     in ls ++ [curr]
   where
     -- TODO: In order to optimize this, use a difference list instead of a
     -- regular list for 'ls'.
     step (ls, curr, width) str
-        | width' > maxWidth = (ls ++ [curr], spaces ++ str, indent + len)
+        | width' > maxWidth = (ls ++ [curr], indent ind str, ind + len)
         | otherwise         = (ls, curr ++ " " ++ str, width')
       where
         len    = length str
         width' = width + 1 + len
 
-    spaces = replicate indent ' '
+
+--------------------------------------------------------------------------------
+withHead :: (a -> a) -> [a] -> [a]
+withHead _ []       = []
+withHead f (x : xs) = f x : xs
+
+
+--------------------------------------------------------------------------------
+withLast :: (a -> a) -> [a] -> [a]
+withLast _ []       = []
+withLast f (x : []) = [f x]
+withLast f (x : xs) = x : withLast f xs
+
+
+--------------------------------------------------------------------------------
+withInit :: (a -> a) -> [a] -> [a]
+withInit _ []       = []
+withInit _ (x : []) = [x]
+withInit f (x : xs) = f x : withInit f xs