🚑 [serokell#45] Save comments inside imports

src/Importify/Pretty.hs

@@ -1,6 +1,5 @@
 -- | This module contains functions to print imports
 -- preserving original formatting.
-
 module Importify.Pretty
        ( printLovelyImports
        ) where
@@ -10,6 +9,7 @@ import           Universum
 import           Control.Arrow         ((&&&))
 import qualified Data.IntMap.Strict    as IM
 import           Data.Text             (strip)
+import qualified Data.Text as T
 import           Language.Haskell.Exts (ImportDecl (..), SrcSpanInfo (..), exactPrint,
                                         srcSpanStartLine)
 
@@ -28,19 +28,20 @@ printLovelyImports start end importsText importDecls =
         indexedImportLines = map (importStartLine &&& exactPrintImport) importDecls
         importsMap         = IM.fromList indexedImportLines
 
-        -- find indexes of empty lines
+        -- find empty lines
         indexedTextLines  = zip [start..end] importsText
-        emptyLinesIndexes = map fst
-                          $ filter (null . strip . stripEndLineComment . snd)
-                                   indexedTextLines
+        emptyLines        = filter (null . strip . snd) indexedTextLines
+
+        -- find single line comments among imports
+        commentLines = filter (T.isPrefixOf "--" . strip . snd) indexedTextLines
 
-        -- add empty lines to result map
-        importsMapWithEmptyLines = foldl' (\dict i -> IM.insert i [""] dict)
+        -- add empty lines and comments to result map
+        importsMapWithExtraLines = foldl' (\dict (i,l) -> IM.insert i [l] dict)
                                           importsMap
-                                          emptyLinesIndexes
+                                          (emptyLines ++ commentLines)
 
         -- collect all values and concat them; order is guaranteed by IntMap
-        resultLines = concat $ IM.elems importsMapWithEmptyLines
+        resultLines = concat $ IM.elems importsMapWithExtraLines
     in resultLines
 
 importStartLine :: ImportDecl SrcSpanInfo -> Int