Allow disabling of Ormolu with special comments

Ormolu can be turned on and off via the special comments:

{- ORMOLU_DISABLE -}

and

{- ORMOLU_ENABLE -}

This allows us to disable formatting selectively for code between these
markers or disable it for the entire file. To achieve the latter, just put
{- ORMOLU_DISABLE -} at the very top. Note that the source code should still
be parseable even without the “excluded” part. Because of that the magic
comments cannot be placed arbitrary, but should rather enclose independent
top-level definitions.

CHANGELOG.md

@@ -9,6 +9,9 @@
 * Comments on pragmas are now preserved. [Issue
   216](https://github.com/tweag/ormolu/issues/216).
 
+* Ormolu can now be enabled and disabled via special comments. [Issue
+  435](https://github.com/tweag/ormolu/issues/435).
+
 ## Ormolu 0.0.4.0
 
 * When given several files to format, Ormolu does not stop on the first

README.md

@@ -91,6 +91,27 @@ formatted output.
 $ ormolu --mode inplace Module.hs
 ```
 
+## Magic comments
+
+Ormolu understands two magic comments:
+
+```haskell
+{- ORMOLU_DISABLE -}
+```
+
+and
+
+```haskell
+{- ORMOLU_ENABLE -}
+```
+
+This allows us to disable formatting selectively for code between these
+markers or disable it for the entire file. To achieve the latter, just put
+`{- ORMOLU_DISABLE -}` at the very top. Note that the source code should
+still be parseable even without the “excluded” part. Because of that the
+magic comments cannot be placed arbitrary, but should rather enclose
+independent top-level definitions.
+
 ## Current limitations
 
 * Does not handle CPP (wontfix, see [the design document][design]).

data/examples/other/disabling/single-definition-out.hs

@@ -0,0 +1,10 @@
+module Foo (foo, bar) where
+
+{- ORMOLU_DISABLE -}
+foo :: Int -> Int
+foo = (+5)
+{- ORMOLU_ENABLE -}
+
+bar :: Bool -> Bool
+bar True = True
+bar False = True

data/examples/other/disabling/single-definition.hs

@@ -0,0 +1,10 @@
+module Foo (foo,bar) where
+
+{- ORMOLU_DISABLE -}
+foo :: Int -> Int
+foo = (+5)
+{- ORMOLU_ENABLE -}
+
+bar :: Bool -> Bool
+bar True  = True
+bar False = True

data/examples/other/disabling/whole-file-out.hs

@@ -0,0 +1,10 @@
+{- ORMOLU_DISABLE -}
+
+module Foo (foo,bar) where
+
+foo :: Int -> Int
+foo = (+5)
+
+bar :: Bool -> Bool
+bar True  = True
+bar False = True

data/examples/other/disabling/whole-file.hs

@@ -0,0 +1,10 @@
+{-    ORMOLU_DISABLE     -}
+
+module Foo (foo,bar) where
+
+foo :: Int -> Int
+foo = (+5)
+
+bar :: Bool -> Bool
+bar True  = True
+bar False = True

ormolu.cabal

@@ -109,6 +109,9 @@ library
                     , Ormolu.Printer.Meat.Type
                     , Ormolu.Printer.Operators
                     , Ormolu.Printer.SpanStream
+                    , Ormolu.Processing.Common
+                    , Ormolu.Processing.Postprocess
+                    , Ormolu.Processing.Preprocess
                     , Ormolu.Utils
   other-modules:      GHC
                     , GHC.DynFlags

src/Ormolu/Parser.hs

@@ -14,9 +14,8 @@ import qualified CmdLineParser as GHC
 import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class
-import Data.List ((\\), foldl', isPrefixOf, sortOn)
+import qualified Data.List as L
 import qualified Data.List.NonEmpty as NE
-import Data.Maybe (catMaybes)
 import Data.Ord (Down (Down))
 import DynFlags as GHC
 import ErrUtils (Severity (..), errMsgSeverity, errMsgSpan)
@@ -32,7 +31,7 @@ import Ormolu.Exception
 import Ormolu.Parser.Anns
 import Ormolu.Parser.CommentStream
 import Ormolu.Parser.Result
-import Ormolu.Parser.Shebang
+import Ormolu.Processing.Preprocess (preprocess)
 import qualified Panic as GHC
 import qualified Parser as GHC
 import qualified StringBuffer as GHC
@@ -51,7 +50,7 @@ parseModule ::
       Either (SrcSpan, String) ParseResult
     )
 parseModule Config {..} path input' = liftIO $ do
-  let (input, extraComments) = extractCommentsFromLines path input'
+  let (input, extraComments) = preprocess path input'
   -- It's important that 'setDefaultExts' is done before
   -- 'parsePragmasIntoDynFlags', because otherwise we might enable an
   -- extension that was explicitly disabled in the file.
@@ -78,7 +77,7 @@ parseModule Config {..} path input' = liftIO $ do
             (pluginModNames dynFlags)
       pStateErrors = \pstate ->
         let errs = bagToList $ GHC.getErrorMessages pstate dynFlags
-         in case sortOn (Down . SeverityOrd . errMsgSeverity) errs of
+         in case L.sortOn (Down . SeverityOrd . errMsgSeverity) errs of
               [] -> Nothing
               err : _ -> Just (errMsgSpan err, show err) -- Show instance returns a short error message
       r = case runParser GHC.parseModule dynFlags path input of
@@ -110,6 +109,14 @@ parseModule Config {..} path input' = liftIO $ do
                       }
   return (warnings, r)
 
+-- | Enable all language extensions that we think should be enabled by
+-- default for ease of use.
+setDefaultExts :: DynFlags -> DynFlags
+setDefaultExts flags = L.foldl' xopt_set flags autoExts
+  where
+    autoExts = allExts L.\\ manualExts
+    allExts = [minBound .. maxBound]
+
 -- | Extensions that are not enabled automatically and should be activated
 -- by user.
 manualExts :: [Extension]
@@ -136,9 +143,6 @@ manualExts =
     -- decision of enabling this style is left to the user
   ]
 
-----------------------------------------------------------------------------
--- Helpers (taken from ghc-exactprint)
-
 -- | Run a 'GHC.P' computation.
 runParser ::
   -- | Computation to run
@@ -157,67 +161,27 @@ runParser parser flags filename input = GHC.unP parser parseState
     buffer = GHC.stringToStringBuffer input
     parseState = GHC.mkPState flags buffer location
 
--- | Transform given input possibly returning comments extracted from it.
--- This handles LINE pragmas and shebangs.
-extractCommentsFromLines ::
-  -- | File name, just to use in the spans
-  FilePath ->
-  -- | Contents of that file
-  String ->
-  -- | Adjusted input with comments extracted from it
-  (String, [Located String])
-extractCommentsFromLines path =
-  unlines' . unzip . zipWith (extractCommentFromLine path) [1 ..] . lines
-  where
-    unlines' (a, b) = (unlines a, catMaybes b)
+-- | Wrap GHC's 'Severity' to add 'Ord' instance.
+newtype SeverityOrd = SeverityOrd Severity
 
--- | Transform a given line possibly returning a comment extracted from it.
-extractCommentFromLine ::
-  -- | File name, just to use in the spans
-  FilePath ->
-  -- | Line number of this line
-  Int ->
-  -- | The actual line
-  String ->
-  -- | Adjusted line and possibly a comment extracted from it
-  (String, Maybe (Located String))
-extractCommentFromLine path line s
-  | "{-# LINE" `isPrefixOf` s =
-    let (pragma, res) = getPragma s
-        size = length pragma
-        ss = mkSrcSpan (mkSrcLoc' 1) (mkSrcLoc' (size + 1))
-     in (res, Just $ L ss pragma)
-  | isShebang s =
-    let ss = mkSrcSpan (mkSrcLoc' 1) (mkSrcLoc' (length s))
-     in ("", Just $ L ss s)
-  | otherwise = (s, Nothing)
-  where
-    mkSrcLoc' = mkSrcLoc (GHC.mkFastString path) line
+instance Eq SeverityOrd where
+  s1 == s2 = compare s1 s2 == EQ
 
--- | Take a line pragma and output its replacement (where line pragma is
--- replaced with spaces) and the contents of the pragma itself.
-getPragma ::
-  -- | Pragma line to analyze
-  String ->
-  -- | Contents of the pragma and its replacement line
-  (String, String)
-getPragma [] = error "Ormolu.Parser.getPragma: input must not be empty"
-getPragma s@(x : xs)
-  | "#-}" `isPrefixOf` s = ("#-}", "   " ++ drop 3 s)
-  | otherwise =
-    let (prag, remline) = getPragma xs
-     in (x : prag, ' ' : remline)
-
--- | Enable all language extensions that we think should be enabled by
--- default for ease of use.
-setDefaultExts :: DynFlags -> DynFlags
-setDefaultExts flags = foldl' GHC.xopt_set flags autoExts
-  where
-    autoExts = allExts \\ manualExts
-    allExts = [minBound .. maxBound]
+instance Ord SeverityOrd where
+  compare (SeverityOrd s1) (SeverityOrd s2) =
+    compare (f s1) (f s2)
+    where
+      f :: Severity -> Int
+      f SevOutput = 1
+      f SevFatal = 2
+      f SevInteractive = 3
+      f SevDump = 4
+      f SevInfo = 5
+      f SevWarning = 6
+      f SevError = 7
 
 ----------------------------------------------------------------------------
--- More helpers (taken from HLint)
+-- Helpers taken from HLint
 
 parsePragmasIntoDynFlags ::
   -- | Pre-set 'DynFlags'
@@ -246,25 +210,3 @@ parsePragmasIntoDynFlags flags extraOpts filepath str =
         reportErr
         (GHC.handleSourceError reportErr act)
     reportErr e = return $ Left (show e)
-
-----------------------------------------------------------------------------
--- Even more helpers
-
--- Wrap GHC's ErrUtils.Severity to add Ord instance
-newtype SeverityOrd = SeverityOrd Severity
-
-instance Eq SeverityOrd where
-  s1 == s2 = compare s1 s2 == EQ
-
-instance Ord SeverityOrd where
-  compare (SeverityOrd s1) (SeverityOrd s2) =
-    compare (f s1) (f s2)
-    where
-      f :: Severity -> Int
-      f SevOutput = 1
-      f SevFatal = 2
-      f SevInteractive = 3
-      f SevDump = 4
-      f SevInfo = 5
-      f SevWarning = 6
-      f SevError = 7

src/Ormolu/Printer.hs

@@ -11,6 +11,7 @@ import Ormolu.Parser.Result
 import Ormolu.Printer.Combinators
 import Ormolu.Printer.Meat.Module
 import Ormolu.Printer.SpanStream
+import Ormolu.Processing.Postprocess (postprocess)
 
 -- | Render a module.
 printModule ::
@@ -19,15 +20,16 @@ printModule ::
   -- | Resulting rendition
   Text
 printModule ParseResult {..} =
-  runR
-    ( p_hsModule
-        prStackHeader
-        prShebangs
-        prPragmas
-        prImportQualifiedPost
-        prParsedSource
-    )
-    (mkSpanStream prParsedSource)
-    prCommentStream
-    prAnns
-    prUseRecordDot
+  postprocess $
+    runR
+      ( p_hsModule
+          prStackHeader
+          prShebangs
+          prPragmas
+          prImportQualifiedPost
+          prParsedSource
+      )
+      (mkSpanStream prParsedSource)
+      prCommentStream
+      prAnns
+      prUseRecordDot

src/Ormolu/Processing/Common.hs

@@ -0,0 +1,27 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+-- | Common definitions for pre- and post- processing.
+module Ormolu.Processing.Common
+  ( OrmoluState (..),
+    startDisabling,
+    endDisabling,
+  )
+where
+
+import Data.String (IsString (..))
+
+-- | Ormolu state.
+data OrmoluState
+  = -- | Enabled
+    OrmoluEnabled
+  | -- | Disabled
+    OrmoluDisabled
+  deriving (Eq, Show)
+
+-- | Marker for the beginning of the region where Ormolu should be disabled.
+startDisabling :: IsString s => s
+startDisabling = "{- ORMOLU_DISABLING_START"
+
+-- | Marker for the end of the region where Ormolu should be disabled.
+endDisabling :: IsString s => s
+endDisabling = "ORMOLU_DISABLE_END -}"

src/Ormolu/Processing/Postprocess.hs

@@ -0,0 +1,16 @@
+-- | Postprocessing for the results of printing.
+module Ormolu.Processing.Postprocess
+  ( postprocess,
+  )
+where
+
+import Data.Text (Text)
+import qualified Data.Text as T
+import Ormolu.Processing.Common
+
+-- | Postprocess output of the formatter.
+postprocess :: Text -> Text
+postprocess = T.unlines . filter (not . magicComment) . T.lines
+  where
+    magicComment x =
+      x == startDisabling || x == endDisabling

src/Ormolu/Processing/Preprocess.hs

@@ -0,0 +1,128 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE LambdaCase #-}
+
+-- | Preprocessing for input source code.
+module Ormolu.Processing.Preprocess
+  ( preprocess,
+  )
+where
+
+import Control.Monad
+import Data.Char (isSpace)
+import qualified Data.List as L
+import Data.Maybe (isJust)
+import Data.Maybe (maybeToList)
+import FastString
+import Ormolu.Parser.Shebang (isShebang)
+import Ormolu.Processing.Common
+import SrcLoc
+
+-- | Transform given input possibly returning comments extracted from it.
+-- This handles LINE pragmas, shebangs, and the magic comments for
+-- enabling\/disabling of Ormolu.
+preprocess ::
+  -- | File name, just to use in the spans
+  FilePath ->
+  -- | Input to process
+  String ->
+  -- | Adjusted input with comments extracted from it
+  (String, [Located String])
+preprocess path input = go 1 OrmoluEnabled id id (lines input)
+  where
+    go !n ormoluState inputSoFar csSoFar = \case
+      [] ->
+        let input' = unlines (inputSoFar [])
+         in ( case ormoluState of
+                OrmoluEnabled -> input'
+                OrmoluDisabled -> input' ++ endDisabling,
+              csSoFar []
+            )
+      (x : xs) ->
+        let (x', ormoluState', cs) = processLine path n ormoluState x
+         in go
+              (n + 1)
+              ormoluState'
+              (inputSoFar . (x' :))
+              (csSoFar . (maybeToList cs ++))
+              xs
+
+-- | Transform a given line possibly returning a comment extracted from it.
+processLine ::
+  -- | File name, just to use in the spans
+  FilePath ->
+  -- | Line number of this line
+  Int ->
+  -- | Whether Ormolu is currently enabled
+  OrmoluState ->
+  -- | The actual line
+  String ->
+  -- | Adjusted line and possibly a comment extracted from it
+  (String, OrmoluState, Maybe (Located String))
+processLine path n ormoluState line
+  | "{-# LINE" `L.isPrefixOf` line =
+    let (pragma, res) = getPragma line
+        size = length pragma
+        ss = mkSrcSpan (mkSrcLoc' 1) (mkSrcLoc' (size + 1))
+     in (res, ormoluState, Just (L ss pragma))
+  | isOrmoluEnable line =
+    case ormoluState of
+      OrmoluEnabled ->
+        (enableMarker, OrmoluEnabled, Nothing)
+      OrmoluDisabled ->
+        (endDisabling ++ enableMarker, OrmoluEnabled, Nothing)
+  | isOrmoluDisable line =
+    case ormoluState of
+      OrmoluEnabled ->
+        (disableMarker ++ startDisabling, OrmoluDisabled, Nothing)
+      OrmoluDisabled ->
+        (disableMarker, OrmoluDisabled, Nothing)
+  | isShebang line =
+    let ss = mkSrcSpan (mkSrcLoc' 1) (mkSrcLoc' (length line))
+     in ("", ormoluState, Just (L ss line))
+  | otherwise = (line, ormoluState, Nothing)
+  where
+    mkSrcLoc' = mkSrcLoc (mkFastString path) n
+
+-- | Take a line pragma and output its replacement (where line pragma is
+-- replaced with spaces) and the contents of the pragma itself.
+getPragma ::
+  -- | Pragma line to analyze
+  String ->
+  -- | Contents of the pragma and its replacement line
+  (String, String)
+getPragma [] = error "Ormolu.Preprocess.getPragma: input must not be empty"
+getPragma s@(x : xs)
+  | "#-}" `L.isPrefixOf` s = ("#-}", "   " ++ drop 3 s)
+  | otherwise =
+    let (prag, remline) = getPragma xs
+     in (x : prag, ' ' : remline)
+
+-- | Canonical enable marker.
+enableMarker :: String
+enableMarker = "{- ORMOLU_ENABLE -}"
+
+-- | Canonical disable marker.
+disableMarker :: String
+disableMarker = "{- ORMOLU_DISABLE -}"
+
+-- | Return 'True' if the given string is an enabling marker.
+isOrmoluEnable :: String -> Bool
+isOrmoluEnable = magicComment "ORMOLU_ENABLE"
+
+-- | Return 'True' if the given string is a disabling marker.
+isOrmoluDisable :: String -> Bool
+isOrmoluDisable = magicComment "ORMOLU_DISABLE"
+
+-- | Construct a function for whitespace-insensitive matching of string.
+magicComment ::
+  -- | What to expect
+  String ->
+  -- | String to test
+  String ->
+  -- | Whether or not the two strings watch
+  Bool
+magicComment expected s0 = isJust $ do
+  s1 <- dropWhile isSpace <$> L.stripPrefix "{-" s0
+  s2 <- dropWhile isSpace <$> L.stripPrefix expected s1
+  s3 <- L.stripPrefix "-}" s2
+  guard (all isSpace s3)