1+ {-# LANGUAGE DerivingStrategies #-}
2+ {-# LANGUAGE GeneralizedNewtypeDeriving #-}
3+ {-# LANGUAGE MultiWayIf #-}
4+ {-# LANGUAGE NamedFieldPuns #-}
5+ {-# LANGUAGE ParallelListComp #-}
6+ {-# LANGUAGE OverloadedStrings #-}
7+ {-# LANGUAGE ViewPatterns #-}
8+
9+ -- | Update YAML preserving top-level key order, blank lines and comments.
10+ --
11+ -- The call sequence is mkRaw, encodeInOrder, redress and unmkRaw but if you
12+ -- don't care about preserving trailing blank lines this can be simplified to
13+ -- encodeInOrder and redress.
14+ --
15+ -- Use yamlLines to transform 'RawYaml' to ['RawYamlLine'].
16+ module Stack.YamlUpdate
17+ ( encodeInOrder , redress
18+ , mkRaw , unmkRaw
19+ , yamlLines
20+ , RawYaml (.. ), RawYamlLine (.. ), YamlKey (.. )
21+ ) where
22+
23+ import Stack.Prelude
24+ import Data.Coerce (coerce )
25+ import qualified Data.List as L
26+ import qualified Data.Yaml as Yaml
27+ import qualified Data.Yaml.Pretty as Yaml
28+ import qualified RIO.Text as T
29+ import qualified RIO.Map as Map
30+
31+ -- | A whole YAML document, may contain line breaks.
32+ newtype RawYaml = RawYaml Text deriving newtype Display
33+ -- | One line from a YAML document, shouldn't contain line breaks.
34+ newtype RawYamlLine = RawYamlLine Text
35+ -- | A YAML top-level key as in @key: value@.
36+ newtype YamlKey = YamlKey Text deriving newtype (Eq , Display )
37+
38+ -- | The line number of a blank line.
39+ newtype YamlLineBlank = YamlLineBlank Int deriving newtype Display
40+ -- | A line number and some content, usually a comment. This can be used with an
41+ -- empty comment to carry the line number for a blank line.
42+ newtype YamlLineComment = YamlLineComment (Int Text )
43+ -- | A mapping from the line number after an encoding that strips blank lines
44+ -- and comments to a line number of the original document.
45+ newtype YamlLineReindex = YamlLineReindex (Int Int
46+
47+ data YamlLines =
48+ YamlLines
49+ { blanks :: ! [YamlLineBlank ]
50+ -- ^ The line numbers of blank lines.
51+ , wholeLineComments :: ! [YamlLineComment ]
52+ -- ^ Comments where # is the first non-space character in that line so
53+ -- that the comment takes up the whole line. Captured with the leading
54+ -- spaces.
55+ , partLineComments :: ! [YamlLineComment ]
56+ -- ^ Comments that have been apended to a line.
57+ , reindices :: ! [YamlLineReindex ]
58+ -- ^ Bumps for line numbers that will need to be moved when blank lines
59+ -- and whole line comments are added back in.
60+ }
61+
62+ data Pegged =
63+ Pegged
64+ { newIndex :: ! Int
65+ -- ^ The new line number to put a line of content.
66+ , leading :: ! [YamlLineComment ]
67+ -- ^ Comments for putting before anything else.
68+ , partComments :: ! [YamlLineComment ]
69+ -- ^ Comments to be appended to lines.
70+ , spanComments :: ! [YamlLineComment ]
71+ -- ^ Blank lines and whole line comments from a range to be put back on
72+ -- the same line as they were taken from.
73+ }
74+
75+ -- | Converts raw YAML as 'Text' with line breaks into a list of lines, dropping
76+ -- trailing line breaks.
77+ yamlLines :: RawYaml -> [RawYamlLine ]
78+ yamlLines x = RawYamlLine <$> T. lines (coerce x)
79+
80+ -- | Puts blank lines and comments from the original lines into the update.
81+ redress :: [RawYamlLine ] -> RawYaml -> RawYaml
82+ redress rawLines (RawYaml t) = let xs = zip [1 .. ] (T. lines t) in RawYaml . T. concat $
83+ [
84+ T. unlines . fromMaybe [x] $ do
85+ Pegged {newIndex = i', leading, partComments, spanComments} <- fetchPegged rawLines (i, j)
86+
87+ let x' = maybe
88+ x
89+ (\ (YamlLineComment (_, c)) -> x <> " " <> dropToComment c)
90+ (L. find ((== i') . commentLineNumber) partComments)
91+
92+ let cs = x' : (comment <$> spanComments)
93+
94+ return $ if i /= 1 then cs else (comment <$> leading) ++ cs
95+
96+ | (i, x) <- xs
97+ | (j, _) <- drop 1 xs ++ [(0 , " "
98+ ]
99+
100+ fetchPegged :: [RawYamlLine ] -> (Int Int -> Maybe Pegged
101+ fetchPegged (pegLines -> yl@ YamlLines {reindices}) (i, j) = do
102+ let reindex = flip L. lookup (coerce reindices)
103+
104+ i' <- reindex i
105+ j' <- reindex j
106+
107+ let (ps, spanned) = fetchInRange yl (\ b -> i' <= b && b < j')
108+
109+ return $ Pegged
110+ { newIndex = i'
111+ , leading = if i /= 1 then [] else snd $ fetchInRange yl (\ b -> b < i')
112+ , partComments = ps
113+ , spanComments = spanned
114+ }
115+
116+ fetchInRange :: YamlLines -> (Int -> Bool -> ([YamlLineComment ], [YamlLineComment ])
117+ fetchInRange YamlLines {blanks, wholeLineComments, partLineComments} p =
118+ let lineNumbers = filter p $ coerce blanks
119+ ls = (\ line -> YamlLineComment (line, " " <$> lineNumbers
120+ filterLineNumber = filter (p . commentLineNumber)
121+ cs = filterLineNumber wholeLineComments
122+ ps = filterLineNumber partLineComments
123+ in
124+ (ps, L. sortOn commentLineNumber $ ls ++ cs)
125+
126+ -- | Uses the order of the keys in the original to preserve the order in the
127+ -- update except that inserting a key orders it last.
128+ encodeInOrder
129+ :: [RawYamlLine ]
130+ -> [YamlKey ]
131+ -> YamlKey
132+ -> Yaml. Object
133+ -> Either UnicodeException RawYaml
134+ encodeInOrder rawLines keysFound upsertKey@ (YamlKey k) yObject =
135+ let keyLine = findKeyLine rawLines
136+ ixMap = Map. fromList $ (\ yk@ (YamlKey x) -> (x, keyLine yk)) <$> keysFound
137+ preservingCompare x y =
138+ -- If updating then preserve order but if inserting then put last.
139+ if | upsertKey `L.elem` keysFound -> Map. lookup x ixMap `compare` Map. lookup y ixMap
140+ | k == x, k == y -> EQ
141+ | k == x -> GT
142+ | k == y -> LT
143+ | otherwise -> Map. lookup x ixMap `compare` Map. lookup y ixMap
144+
145+ keyCmp = Yaml. setConfCompare preservingCompare Yaml. defConfig
146+
147+ in RawYaml <$> decodeUtf8' (Yaml. encodePretty keyCmp yObject)
148+
149+ endSentinel :: Text
150+ endSentinel = " ED10F56C-562E-4847-A50B-7541C1732A15: 2986F150-E4A0-41D8-AB9C-8BD82FA12DC4"
151+
152+ mkRaw :: Text -> RawYaml
153+ mkRaw = addSentinels . RawYaml
154+
155+ unmkRaw :: RawYaml -> Text
156+ unmkRaw = coerce . removeSentinels
157+
158+ -- | This is leaking implementation but adding a sentinal key-value to the end
159+ -- of YAML is a cheap way to ensure trailing newlines are not swallowed.
160+ addSentinels :: RawYaml -> RawYaml
161+ addSentinels (RawYaml x) = RawYaml $ x <> endSentinel
162+
163+ removeSentinels :: RawYaml -> RawYaml
164+ removeSentinels (RawYaml x) = RawYaml . T. unlines . filter (/= endSentinel) $ T. lines x
165+
166+ findKeyLine :: [RawYamlLine ] -> YamlKey -> Maybe Int
167+ findKeyLine rawLines (YamlKey x) = join . listToMaybe . take 1 . dropWhile isNothing $
168+ [ if x `T.isPrefixOf` y then Just i else Nothing
169+ | RawYamlLine y <- rawLines
170+ | i <- [1 .. ]
171+ ]
172+
173+ comment :: YamlLineComment -> Text
174+ comment (YamlLineComment (_, c)) = c
175+
176+ commentLineNumber :: YamlLineComment -> Int
177+ commentLineNumber (YamlLineComment (c, _)) = c
178+
179+ instance Display YamlLineComment where
180+ textDisplay (YamlLineComment (i, s)) = textDisplay . T. pack $ show (i, T. unpack s)
181+
182+ dropToComment :: Text -> Text
183+ dropToComment = T. dropWhile (/= ' #'
184+
185+ -- | Gather enough information about lines to peg line numbers so that blank
186+ -- lines and comments can be reinserted later.
187+ pegLines :: [RawYamlLine ] -> YamlLines
188+ pegLines rawLines =
189+ let (ls, rs) = partitionEithers
190+ [
191+ if | y == " " -> Left . Left $ YamlLineBlank i
192+
193+ | " #" `T.isPrefixOf` T. dropWhile (== ' ' ->
194+ Left . Right $ YamlLineComment (i, y)
195+
196+ | otherwise ->
197+ if " #" `T.isPrefixOf` dropToComment y
198+ then Right . Left $ YamlLineComment (i, y)
199+ else Right $ Right
200+
201+ | RawYamlLine y <- rawLines
202+ | i <- [1 .. ]
203+ ]
204+
205+ (blanks, wholeLineComments) = partitionEithers ls
206+ (partLineComments, contentLines) = partitionEithers rs
207+ indexLines = L. sort $ contentLines ++ (commentLineNumber <$> partLineComments)
208+ reindex = zipWith (curry YamlLineReindex ) [1 .. ] indexLines
209+
210+ in YamlLines blanks wholeLineComments partLineComments reindex
0 commit comments