]> git.madduck.net Git - etc/vim.git/blob - src/black/trans.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

Cover more in the usage docs (#2208)
[etc/vim.git] / src / black / trans.py
1 """
2 String transformers that can split and merge strings.
3 """
4 from abc import ABC, abstractmethod
5 from collections import defaultdict
6 from dataclasses import dataclass
7 import regex as re
8 from typing import (
9     Any,
10     Callable,
11     Collection,
12     Dict,
13     Iterable,
14     Iterator,
15     List,
16     Optional,
17     Sequence,
18     Tuple,
19     TypeVar,
20     Union,
21 )
22
23 from black.rusty import Result, Ok, Err
24
25 from black.mode import Feature
26 from black.nodes import syms, replace_child, parent_type
27 from black.nodes import is_empty_par, is_empty_lpar, is_empty_rpar
28 from black.nodes import CLOSING_BRACKETS, STANDALONE_COMMENT
29 from black.lines import Line, append_leaves
30 from black.brackets import BracketMatchError
31 from black.comments import contains_pragma_comment
32 from black.strings import has_triple_quotes, get_string_prefix, assert_is_leaf_string
33 from black.strings import normalize_string_quotes
34
35 from blib2to3.pytree import Leaf, Node
36 from blib2to3.pgen2 import token
37
38
39 class CannotTransform(Exception):
40     """Base class for errors raised by Transformers."""
41
42
43 # types
44 T = TypeVar("T")
45 LN = Union[Leaf, Node]
46 Transformer = Callable[[Line, Collection[Feature]], Iterator[Line]]
47 Index = int
48 NodeType = int
49 ParserState = int
50 StringID = int
51 TResult = Result[T, CannotTransform]  # (T)ransform Result
52 TMatchResult = TResult[Index]
53
54
55 def TErr(err_msg: str) -> Err[CannotTransform]:
56     """(T)ransform Err
57
58     Convenience function used when working with the TResult type.
59     """
60     cant_transform = CannotTransform(err_msg)
61     return Err(cant_transform)
62
63
64 @dataclass  # type: ignore
65 class StringTransformer(ABC):
66     """
67     An implementation of the Transformer protocol that relies on its
68     subclasses overriding the template methods `do_match(...)` and
69     `do_transform(...)`.
70
71     This Transformer works exclusively on strings (for example, by merging
72     or splitting them).
73
74     The following sections can be found among the docstrings of each concrete
75     StringTransformer subclass.
76
77     Requirements:
78         Which requirements must be met of the given Line for this
79         StringTransformer to be applied?
80
81     Transformations:
82         If the given Line meets all of the above requirements, which string
83         transformations can you expect to be applied to it by this
84         StringTransformer?
85
86     Collaborations:
87         What contractual agreements does this StringTransformer have with other
88         StringTransfomers? Such collaborations should be eliminated/minimized
89         as much as possible.
90     """
91
92     line_length: int
93     normalize_strings: bool
94     __name__ = "StringTransformer"
95
96     @abstractmethod
97     def do_match(self, line: Line) -> TMatchResult:
98         """
99         Returns:
100             * Ok(string_idx) such that `line.leaves[string_idx]` is our target
101             string, if a match was able to be made.
102                 OR
103             * Err(CannotTransform), if a match was not able to be made.
104         """
105
106     @abstractmethod
107     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
108         """
109         Yields:
110             * Ok(new_line) where new_line is the new transformed line.
111                 OR
112             * Err(CannotTransform) if the transformation failed for some reason. The
113             `do_match(...)` template method should usually be used to reject
114             the form of the given Line, but in some cases it is difficult to
115             know whether or not a Line meets the StringTransformer's
116             requirements until the transformation is already midway.
117
118         Side Effects:
119             This method should NOT mutate @line directly, but it MAY mutate the
120             Line's underlying Node structure. (WARNING: If the underlying Node
121             structure IS altered, then this method should NOT be allowed to
122             yield an CannotTransform after that point.)
123         """
124
125     def __call__(self, line: Line, _features: Collection[Feature]) -> Iterator[Line]:
126         """
127         StringTransformer instances have a call signature that mirrors that of
128         the Transformer type.
129
130         Raises:
131             CannotTransform(...) if the concrete StringTransformer class is unable
132             to transform @line.
133         """
134         # Optimization to avoid calling `self.do_match(...)` when the line does
135         # not contain any string.
136         if not any(leaf.type == token.STRING for leaf in line.leaves):
137             raise CannotTransform("There are no strings in this line.")
138
139         match_result = self.do_match(line)
140
141         if isinstance(match_result, Err):
142             cant_transform = match_result.err()
143             raise CannotTransform(
144                 f"The string transformer {self.__class__.__name__} does not recognize"
145                 " this line as one that it can transform."
146             ) from cant_transform
147
148         string_idx = match_result.ok()
149
150         for line_result in self.do_transform(line, string_idx):
151             if isinstance(line_result, Err):
152                 cant_transform = line_result.err()
153                 raise CannotTransform(
154                     "StringTransformer failed while attempting to transform string."
155                 ) from cant_transform
156             line = line_result.ok()
157             yield line
158
159
160 @dataclass
161 class CustomSplit:
162     """A custom (i.e. manual) string split.
163
164     A single CustomSplit instance represents a single substring.
165
166     Examples:
167         Consider the following string:
168         ```
169         "Hi there friend."
170         " This is a custom"
171         f" string {split}."
172         ```
173
174         This string will correspond to the following three CustomSplit instances:
175         ```
176         CustomSplit(False, 16)
177         CustomSplit(False, 17)
178         CustomSplit(True, 16)
179         ```
180     """
181
182     has_prefix: bool
183     break_idx: int
184
185
186 class CustomSplitMapMixin:
187     """
188     This mixin class is used to map merged strings to a sequence of
189     CustomSplits, which will then be used to re-split the strings iff none of
190     the resultant substrings go over the configured max line length.
191     """
192
193     _Key = Tuple[StringID, str]
194     _CUSTOM_SPLIT_MAP: Dict[_Key, Tuple[CustomSplit, ...]] = defaultdict(tuple)
195
196     @staticmethod
197     def _get_key(string: str) -> "CustomSplitMapMixin._Key":
198         """
199         Returns:
200             A unique identifier that is used internally to map @string to a
201             group of custom splits.
202         """
203         return (id(string), string)
204
205     def add_custom_splits(
206         self, string: str, custom_splits: Iterable[CustomSplit]
207     ) -> None:
208         """Custom Split Map Setter Method
209
210         Side Effects:
211             Adds a mapping from @string to the custom splits @custom_splits.
212         """
213         key = self._get_key(string)
214         self._CUSTOM_SPLIT_MAP[key] = tuple(custom_splits)
215
216     def pop_custom_splits(self, string: str) -> List[CustomSplit]:
217         """Custom Split Map Getter Method
218
219         Returns:
220             * A list of the custom splits that are mapped to @string, if any
221             exist.
222                 OR
223             * [], otherwise.
224
225         Side Effects:
226             Deletes the mapping between @string and its associated custom
227             splits (which are returned to the caller).
228         """
229         key = self._get_key(string)
230
231         custom_splits = self._CUSTOM_SPLIT_MAP[key]
232         del self._CUSTOM_SPLIT_MAP[key]
233
234         return list(custom_splits)
235
236     def has_custom_splits(self, string: str) -> bool:
237         """
238         Returns:
239             True iff @string is associated with a set of custom splits.
240         """
241         key = self._get_key(string)
242         return key in self._CUSTOM_SPLIT_MAP
243
244
245 class StringMerger(CustomSplitMapMixin, StringTransformer):
246     """StringTransformer that merges strings together.
247
248     Requirements:
249         (A) The line contains adjacent strings such that ALL of the validation checks
250         listed in StringMerger.__validate_msg(...)'s docstring pass.
251             OR
252         (B) The line contains a string which uses line continuation backslashes.
253
254     Transformations:
255         Depending on which of the two requirements above where met, either:
256
257         (A) The string group associated with the target string is merged.
258             OR
259         (B) All line-continuation backslashes are removed from the target string.
260
261     Collaborations:
262         StringMerger provides custom split information to StringSplitter.
263     """
264
265     def do_match(self, line: Line) -> TMatchResult:
266         LL = line.leaves
267
268         is_valid_index = is_valid_index_factory(LL)
269
270         for (i, leaf) in enumerate(LL):
271             if (
272                 leaf.type == token.STRING
273                 and is_valid_index(i + 1)
274                 and LL[i + 1].type == token.STRING
275             ):
276                 return Ok(i)
277
278             if leaf.type == token.STRING and "\\\n" in leaf.value:
279                 return Ok(i)
280
281         return TErr("This line has no strings that need merging.")
282
283     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
284         new_line = line
285         rblc_result = self._remove_backslash_line_continuation_chars(
286             new_line, string_idx
287         )
288         if isinstance(rblc_result, Ok):
289             new_line = rblc_result.ok()
290
291         msg_result = self._merge_string_group(new_line, string_idx)
292         if isinstance(msg_result, Ok):
293             new_line = msg_result.ok()
294
295         if isinstance(rblc_result, Err) and isinstance(msg_result, Err):
296             msg_cant_transform = msg_result.err()
297             rblc_cant_transform = rblc_result.err()
298             cant_transform = CannotTransform(
299                 "StringMerger failed to merge any strings in this line."
300             )
301
302             # Chain the errors together using `__cause__`.
303             msg_cant_transform.__cause__ = rblc_cant_transform
304             cant_transform.__cause__ = msg_cant_transform
305
306             yield Err(cant_transform)
307         else:
308             yield Ok(new_line)
309
310     @staticmethod
311     def _remove_backslash_line_continuation_chars(
312         line: Line, string_idx: int
313     ) -> TResult[Line]:
314         """
315         Merge strings that were split across multiple lines using
316         line-continuation backslashes.
317
318         Returns:
319             Ok(new_line), if @line contains backslash line-continuation
320             characters.
321                 OR
322             Err(CannotTransform), otherwise.
323         """
324         LL = line.leaves
325
326         string_leaf = LL[string_idx]
327         if not (
328             string_leaf.type == token.STRING
329             and "\\\n" in string_leaf.value
330             and not has_triple_quotes(string_leaf.value)
331         ):
332             return TErr(
333                 f"String leaf {string_leaf} does not contain any backslash line"
334                 " continuation characters."
335             )
336
337         new_line = line.clone()
338         new_line.comments = line.comments.copy()
339         append_leaves(new_line, line, LL)
340
341         new_string_leaf = new_line.leaves[string_idx]
342         new_string_leaf.value = new_string_leaf.value.replace("\\\n", "")
343
344         return Ok(new_line)
345
346     def _merge_string_group(self, line: Line, string_idx: int) -> TResult[Line]:
347         """
348         Merges string group (i.e. set of adjacent strings) where the first
349         string in the group is `line.leaves[string_idx]`.
350
351         Returns:
352             Ok(new_line), if ALL of the validation checks found in
353             __validate_msg(...) pass.
354                 OR
355             Err(CannotTransform), otherwise.
356         """
357         LL = line.leaves
358
359         is_valid_index = is_valid_index_factory(LL)
360
361         vresult = self._validate_msg(line, string_idx)
362         if isinstance(vresult, Err):
363             return vresult
364
365         # If the string group is wrapped inside an Atom node, we must make sure
366         # to later replace that Atom with our new (merged) string leaf.
367         atom_node = LL[string_idx].parent
368
369         # We will place BREAK_MARK in between every two substrings that we
370         # merge. We will then later go through our final result and use the
371         # various instances of BREAK_MARK we find to add the right values to
372         # the custom split map.
373         BREAK_MARK = "@@@@@ BLACK BREAKPOINT MARKER @@@@@"
374
375         QUOTE = LL[string_idx].value[-1]
376
377         def make_naked(string: str, string_prefix: str) -> str:
378             """Strip @string (i.e. make it a "naked" string)
379
380             Pre-conditions:
381                 * assert_is_leaf_string(@string)
382
383             Returns:
384                 A string that is identical to @string except that
385                 @string_prefix has been stripped, the surrounding QUOTE
386                 characters have been removed, and any remaining QUOTE
387                 characters have been escaped.
388             """
389             assert_is_leaf_string(string)
390
391             RE_EVEN_BACKSLASHES = r"(?:(?<!\\)(?:\\\\)*)"
392             naked_string = string[len(string_prefix) + 1 : -1]
393             naked_string = re.sub(
394                 "(" + RE_EVEN_BACKSLASHES + ")" + QUOTE, r"\1\\" + QUOTE, naked_string
395             )
396             return naked_string
397
398         # Holds the CustomSplit objects that will later be added to the custom
399         # split map.
400         custom_splits = []
401
402         # Temporary storage for the 'has_prefix' part of the CustomSplit objects.
403         prefix_tracker = []
404
405         # Sets the 'prefix' variable. This is the prefix that the final merged
406         # string will have.
407         next_str_idx = string_idx
408         prefix = ""
409         while (
410             not prefix
411             and is_valid_index(next_str_idx)
412             and LL[next_str_idx].type == token.STRING
413         ):
414             prefix = get_string_prefix(LL[next_str_idx].value)
415             next_str_idx += 1
416
417         # The next loop merges the string group. The final string will be
418         # contained in 'S'.
419         #
420         # The following convenience variables are used:
421         #
422         #   S: string
423         #   NS: naked string
424         #   SS: next string
425         #   NSS: naked next string
426         S = ""
427         NS = ""
428         num_of_strings = 0
429         next_str_idx = string_idx
430         while is_valid_index(next_str_idx) and LL[next_str_idx].type == token.STRING:
431             num_of_strings += 1
432
433             SS = LL[next_str_idx].value
434             next_prefix = get_string_prefix(SS)
435
436             # If this is an f-string group but this substring is not prefixed
437             # with 'f'...
438             if "f" in prefix and "f" not in next_prefix:
439                 # Then we must escape any braces contained in this substring.
440                 SS = re.subf(r"(\{|\})", "{1}{1}", SS)
441
442             NSS = make_naked(SS, next_prefix)
443
444             has_prefix = bool(next_prefix)
445             prefix_tracker.append(has_prefix)
446
447             S = prefix + QUOTE + NS + NSS + BREAK_MARK + QUOTE
448             NS = make_naked(S, prefix)
449
450             next_str_idx += 1
451
452         S_leaf = Leaf(token.STRING, S)
453         if self.normalize_strings:
454             S_leaf.value = normalize_string_quotes(S_leaf.value)
455
456         # Fill the 'custom_splits' list with the appropriate CustomSplit objects.
457         temp_string = S_leaf.value[len(prefix) + 1 : -1]
458         for has_prefix in prefix_tracker:
459             mark_idx = temp_string.find(BREAK_MARK)
460             assert (
461                 mark_idx >= 0
462             ), "Logic error while filling the custom string breakpoint cache."
463
464             temp_string = temp_string[mark_idx + len(BREAK_MARK) :]
465             breakpoint_idx = mark_idx + (len(prefix) if has_prefix else 0) + 1
466             custom_splits.append(CustomSplit(has_prefix, breakpoint_idx))
467
468         string_leaf = Leaf(token.STRING, S_leaf.value.replace(BREAK_MARK, ""))
469
470         if atom_node is not None:
471             replace_child(atom_node, string_leaf)
472
473         # Build the final line ('new_line') that this method will later return.
474         new_line = line.clone()
475         for (i, leaf) in enumerate(LL):
476             if i == string_idx:
477                 new_line.append(string_leaf)
478
479             if string_idx <= i < string_idx + num_of_strings:
480                 for comment_leaf in line.comments_after(LL[i]):
481                     new_line.append(comment_leaf, preformatted=True)
482                 continue
483
484             append_leaves(new_line, line, [leaf])
485
486         self.add_custom_splits(string_leaf.value, custom_splits)
487         return Ok(new_line)
488
489     @staticmethod
490     def _validate_msg(line: Line, string_idx: int) -> TResult[None]:
491         """Validate (M)erge (S)tring (G)roup
492
493         Transform-time string validation logic for __merge_string_group(...).
494
495         Returns:
496             * Ok(None), if ALL validation checks (listed below) pass.
497                 OR
498             * Err(CannotTransform), if any of the following are true:
499                 - The target string group does not contain ANY stand-alone comments.
500                 - The target string is not in a string group (i.e. it has no
501                   adjacent strings).
502                 - The string group has more than one inline comment.
503                 - The string group has an inline comment that appears to be a pragma.
504                 - The set of all string prefixes in the string group is of
505                   length greater than one and is not equal to {"", "f"}.
506                 - The string group consists of raw strings.
507         """
508         # We first check for "inner" stand-alone comments (i.e. stand-alone
509         # comments that have a string leaf before them AND after them).
510         for inc in [1, -1]:
511             i = string_idx
512             found_sa_comment = False
513             is_valid_index = is_valid_index_factory(line.leaves)
514             while is_valid_index(i) and line.leaves[i].type in [
515                 token.STRING,
516                 STANDALONE_COMMENT,
517             ]:
518                 if line.leaves[i].type == STANDALONE_COMMENT:
519                     found_sa_comment = True
520                 elif found_sa_comment:
521                     return TErr(
522                         "StringMerger does NOT merge string groups which contain "
523                         "stand-alone comments."
524                     )
525
526                 i += inc
527
528         num_of_inline_string_comments = 0
529         set_of_prefixes = set()
530         num_of_strings = 0
531         for leaf in line.leaves[string_idx:]:
532             if leaf.type != token.STRING:
533                 # If the string group is trailed by a comma, we count the
534                 # comments trailing the comma to be one of the string group's
535                 # comments.
536                 if leaf.type == token.COMMA and id(leaf) in line.comments:
537                     num_of_inline_string_comments += 1
538                 break
539
540             if has_triple_quotes(leaf.value):
541                 return TErr("StringMerger does NOT merge multiline strings.")
542
543             num_of_strings += 1
544             prefix = get_string_prefix(leaf.value)
545             if "r" in prefix:
546                 return TErr("StringMerger does NOT merge raw strings.")
547
548             set_of_prefixes.add(prefix)
549
550             if id(leaf) in line.comments:
551                 num_of_inline_string_comments += 1
552                 if contains_pragma_comment(line.comments[id(leaf)]):
553                     return TErr("Cannot merge strings which have pragma comments.")
554
555         if num_of_strings < 2:
556             return TErr(
557                 f"Not enough strings to merge (num_of_strings={num_of_strings})."
558             )
559
560         if num_of_inline_string_comments > 1:
561             return TErr(
562                 f"Too many inline string comments ({num_of_inline_string_comments})."
563             )
564
565         if len(set_of_prefixes) > 1 and set_of_prefixes != {"", "f"}:
566             return TErr(f"Too many different prefixes ({set_of_prefixes}).")
567
568         return Ok(None)
569
570
571 class StringParenStripper(StringTransformer):
572     """StringTransformer that strips surrounding parentheses from strings.
573
574     Requirements:
575         The line contains a string which is surrounded by parentheses and:
576             - The target string is NOT the only argument to a function call.
577             - The target string is NOT a "pointless" string.
578             - If the target string contains a PERCENT, the brackets are not
579               preceeded or followed by an operator with higher precedence than
580               PERCENT.
581
582     Transformations:
583         The parentheses mentioned in the 'Requirements' section are stripped.
584
585     Collaborations:
586         StringParenStripper has its own inherent usefulness, but it is also
587         relied on to clean up the parentheses created by StringParenWrapper (in
588         the event that they are no longer needed).
589     """
590
591     def do_match(self, line: Line) -> TMatchResult:
592         LL = line.leaves
593
594         is_valid_index = is_valid_index_factory(LL)
595
596         for (idx, leaf) in enumerate(LL):
597             # Should be a string...
598             if leaf.type != token.STRING:
599                 continue
600
601             # If this is a "pointless" string...
602             if (
603                 leaf.parent
604                 and leaf.parent.parent
605                 and leaf.parent.parent.type == syms.simple_stmt
606             ):
607                 continue
608
609             # Should be preceded by a non-empty LPAR...
610             if (
611                 not is_valid_index(idx - 1)
612                 or LL[idx - 1].type != token.LPAR
613                 or is_empty_lpar(LL[idx - 1])
614             ):
615                 continue
616
617             # That LPAR should NOT be preceded by a function name or a closing
618             # bracket (which could be a function which returns a function or a
619             # list/dictionary that contains a function)...
620             if is_valid_index(idx - 2) and (
621                 LL[idx - 2].type == token.NAME or LL[idx - 2].type in CLOSING_BRACKETS
622             ):
623                 continue
624
625             string_idx = idx
626
627             # Skip the string trailer, if one exists.
628             string_parser = StringParser()
629             next_idx = string_parser.parse(LL, string_idx)
630
631             # if the leaves in the parsed string include a PERCENT, we need to
632             # make sure the initial LPAR is NOT preceded by an operator with
633             # higher or equal precedence to PERCENT
634             if is_valid_index(idx - 2):
635                 # mypy can't quite follow unless we name this
636                 before_lpar = LL[idx - 2]
637                 if token.PERCENT in {leaf.type for leaf in LL[idx - 1 : next_idx]} and (
638                     (
639                         before_lpar.type
640                         in {
641                             token.STAR,
642                             token.AT,
643                             token.SLASH,
644                             token.DOUBLESLASH,
645                             token.PERCENT,
646                             token.TILDE,
647                             token.DOUBLESTAR,
648                             token.AWAIT,
649                             token.LSQB,
650                             token.LPAR,
651                         }
652                     )
653                     or (
654                         # only unary PLUS/MINUS
655                         before_lpar.parent
656                         and before_lpar.parent.type == syms.factor
657                         and (before_lpar.type in {token.PLUS, token.MINUS})
658                     )
659                 ):
660                     continue
661
662             # Should be followed by a non-empty RPAR...
663             if (
664                 is_valid_index(next_idx)
665                 and LL[next_idx].type == token.RPAR
666                 and not is_empty_rpar(LL[next_idx])
667             ):
668                 # That RPAR should NOT be followed by anything with higher
669                 # precedence than PERCENT
670                 if is_valid_index(next_idx + 1) and LL[next_idx + 1].type in {
671                     token.DOUBLESTAR,
672                     token.LSQB,
673                     token.LPAR,
674                     token.DOT,
675                 }:
676                     continue
677
678                 return Ok(string_idx)
679
680         return TErr("This line has no strings wrapped in parens.")
681
682     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
683         LL = line.leaves
684
685         string_parser = StringParser()
686         rpar_idx = string_parser.parse(LL, string_idx)
687
688         for leaf in (LL[string_idx - 1], LL[rpar_idx]):
689             if line.comments_after(leaf):
690                 yield TErr(
691                     "Will not strip parentheses which have comments attached to them."
692                 )
693                 return
694
695         new_line = line.clone()
696         new_line.comments = line.comments.copy()
697         try:
698             append_leaves(new_line, line, LL[: string_idx - 1])
699         except BracketMatchError:
700             # HACK: I believe there is currently a bug somewhere in
701             # right_hand_split() that is causing brackets to not be tracked
702             # properly by a shared BracketTracker.
703             append_leaves(new_line, line, LL[: string_idx - 1], preformatted=True)
704
705         string_leaf = Leaf(token.STRING, LL[string_idx].value)
706         LL[string_idx - 1].remove()
707         replace_child(LL[string_idx], string_leaf)
708         new_line.append(string_leaf)
709
710         append_leaves(
711             new_line, line, LL[string_idx + 1 : rpar_idx] + LL[rpar_idx + 1 :]
712         )
713
714         LL[rpar_idx].remove()
715
716         yield Ok(new_line)
717
718
719 class BaseStringSplitter(StringTransformer):
720     """
721     Abstract class for StringTransformers which transform a Line's strings by splitting
722     them or placing them on their own lines where necessary to avoid going over
723     the configured line length.
724
725     Requirements:
726         * The target string value is responsible for the line going over the
727         line length limit. It follows that after all of black's other line
728         split methods have been exhausted, this line (or one of the resulting
729         lines after all line splits are performed) would still be over the
730         line_length limit unless we split this string.
731             AND
732         * The target string is NOT a "pointless" string (i.e. a string that has
733         no parent or siblings).
734             AND
735         * The target string is not followed by an inline comment that appears
736         to be a pragma.
737             AND
738         * The target string is not a multiline (i.e. triple-quote) string.
739     """
740
741     @abstractmethod
742     def do_splitter_match(self, line: Line) -> TMatchResult:
743         """
744         BaseStringSplitter asks its clients to override this method instead of
745         `StringTransformer.do_match(...)`.
746
747         Follows the same protocol as `StringTransformer.do_match(...)`.
748
749         Refer to `help(StringTransformer.do_match)` for more information.
750         """
751
752     def do_match(self, line: Line) -> TMatchResult:
753         match_result = self.do_splitter_match(line)
754         if isinstance(match_result, Err):
755             return match_result
756
757         string_idx = match_result.ok()
758         vresult = self._validate(line, string_idx)
759         if isinstance(vresult, Err):
760             return vresult
761
762         return match_result
763
764     def _validate(self, line: Line, string_idx: int) -> TResult[None]:
765         """
766         Checks that @line meets all of the requirements listed in this classes'
767         docstring. Refer to `help(BaseStringSplitter)` for a detailed
768         description of those requirements.
769
770         Returns:
771             * Ok(None), if ALL of the requirements are met.
772                 OR
773             * Err(CannotTransform), if ANY of the requirements are NOT met.
774         """
775         LL = line.leaves
776
777         string_leaf = LL[string_idx]
778
779         max_string_length = self._get_max_string_length(line, string_idx)
780         if len(string_leaf.value) <= max_string_length:
781             return TErr(
782                 "The string itself is not what is causing this line to be too long."
783             )
784
785         if not string_leaf.parent or [L.type for L in string_leaf.parent.children] == [
786             token.STRING,
787             token.NEWLINE,
788         ]:
789             return TErr(
790                 f"This string ({string_leaf.value}) appears to be pointless (i.e. has"
791                 " no parent)."
792             )
793
794         if id(line.leaves[string_idx]) in line.comments and contains_pragma_comment(
795             line.comments[id(line.leaves[string_idx])]
796         ):
797             return TErr(
798                 "Line appears to end with an inline pragma comment. Splitting the line"
799                 " could modify the pragma's behavior."
800             )
801
802         if has_triple_quotes(string_leaf.value):
803             return TErr("We cannot split multiline strings.")
804
805         return Ok(None)
806
807     def _get_max_string_length(self, line: Line, string_idx: int) -> int:
808         """
809         Calculates the max string length used when attempting to determine
810         whether or not the target string is responsible for causing the line to
811         go over the line length limit.
812
813         WARNING: This method is tightly coupled to both StringSplitter and
814         (especially) StringParenWrapper. There is probably a better way to
815         accomplish what is being done here.
816
817         Returns:
818             max_string_length: such that `line.leaves[string_idx].value >
819             max_string_length` implies that the target string IS responsible
820             for causing this line to exceed the line length limit.
821         """
822         LL = line.leaves
823
824         is_valid_index = is_valid_index_factory(LL)
825
826         # We use the shorthand "WMA4" in comments to abbreviate "We must
827         # account for". When giving examples, we use STRING to mean some/any
828         # valid string.
829         #
830         # Finally, we use the following convenience variables:
831         #
832         #   P:  The leaf that is before the target string leaf.
833         #   N:  The leaf that is after the target string leaf.
834         #   NN: The leaf that is after N.
835
836         # WMA4 the whitespace at the beginning of the line.
837         offset = line.depth * 4
838
839         if is_valid_index(string_idx - 1):
840             p_idx = string_idx - 1
841             if (
842                 LL[string_idx - 1].type == token.LPAR
843                 and LL[string_idx - 1].value == ""
844                 and string_idx >= 2
845             ):
846                 # If the previous leaf is an empty LPAR placeholder, we should skip it.
847                 p_idx -= 1
848
849             P = LL[p_idx]
850             if P.type == token.PLUS:
851                 # WMA4 a space and a '+' character (e.g. `+ STRING`).
852                 offset += 2
853
854             if P.type == token.COMMA:
855                 # WMA4 a space, a comma, and a closing bracket [e.g. `), STRING`].
856                 offset += 3
857
858             if P.type in [token.COLON, token.EQUAL, token.NAME]:
859                 # This conditional branch is meant to handle dictionary keys,
860                 # variable assignments, 'return STRING' statement lines, and
861                 # 'else STRING' ternary expression lines.
862
863                 # WMA4 a single space.
864                 offset += 1
865
866                 # WMA4 the lengths of any leaves that came before that space,
867                 # but after any closing bracket before that space.
868                 for leaf in reversed(LL[: p_idx + 1]):
869                     offset += len(str(leaf))
870                     if leaf.type in CLOSING_BRACKETS:
871                         break
872
873         if is_valid_index(string_idx + 1):
874             N = LL[string_idx + 1]
875             if N.type == token.RPAR and N.value == "" and len(LL) > string_idx + 2:
876                 # If the next leaf is an empty RPAR placeholder, we should skip it.
877                 N = LL[string_idx + 2]
878
879             if N.type == token.COMMA:
880                 # WMA4 a single comma at the end of the string (e.g `STRING,`).
881                 offset += 1
882
883             if is_valid_index(string_idx + 2):
884                 NN = LL[string_idx + 2]
885
886                 if N.type == token.DOT and NN.type == token.NAME:
887                     # This conditional branch is meant to handle method calls invoked
888                     # off of a string literal up to and including the LPAR character.
889
890                     # WMA4 the '.' character.
891                     offset += 1
892
893                     if (
894                         is_valid_index(string_idx + 3)
895                         and LL[string_idx + 3].type == token.LPAR
896                     ):
897                         # WMA4 the left parenthesis character.
898                         offset += 1
899
900                     # WMA4 the length of the method's name.
901                     offset += len(NN.value)
902
903         has_comments = False
904         for comment_leaf in line.comments_after(LL[string_idx]):
905             if not has_comments:
906                 has_comments = True
907                 # WMA4 two spaces before the '#' character.
908                 offset += 2
909
910             # WMA4 the length of the inline comment.
911             offset += len(comment_leaf.value)
912
913         max_string_length = self.line_length - offset
914         return max_string_length
915
916
917 class StringSplitter(CustomSplitMapMixin, BaseStringSplitter):
918     """
919     StringTransformer that splits "atom" strings (i.e. strings which exist on
920     lines by themselves).
921
922     Requirements:
923         * The line consists ONLY of a single string (with the exception of a
924         '+' symbol which MAY exist at the start of the line), MAYBE a string
925         trailer, and MAYBE a trailing comma.
926             AND
927         * All of the requirements listed in BaseStringSplitter's docstring.
928
929     Transformations:
930         The string mentioned in the 'Requirements' section is split into as
931         many substrings as necessary to adhere to the configured line length.
932
933         In the final set of substrings, no substring should be smaller than
934         MIN_SUBSTR_SIZE characters.
935
936         The string will ONLY be split on spaces (i.e. each new substring should
937         start with a space). Note that the string will NOT be split on a space
938         which is escaped with a backslash.
939
940         If the string is an f-string, it will NOT be split in the middle of an
941         f-expression (e.g. in f"FooBar: {foo() if x else bar()}", {foo() if x
942         else bar()} is an f-expression).
943
944         If the string that is being split has an associated set of custom split
945         records and those custom splits will NOT result in any line going over
946         the configured line length, those custom splits are used. Otherwise the
947         string is split as late as possible (from left-to-right) while still
948         adhering to the transformation rules listed above.
949
950     Collaborations:
951         StringSplitter relies on StringMerger to construct the appropriate
952         CustomSplit objects and add them to the custom split map.
953     """
954
955     MIN_SUBSTR_SIZE = 6
956     # Matches an "f-expression" (e.g. {var}) that might be found in an f-string.
957     RE_FEXPR = r"""
958     (?<!\{) (?:\{\{)* \{ (?!\{)
959         (?:
960             [^\{\}]
961             | \{\{
962             | \}\}
963             | (?R)
964         )+?
965     (?<!\}) \} (?:\}\})* (?!\})
966     """
967
968     def do_splitter_match(self, line: Line) -> TMatchResult:
969         LL = line.leaves
970
971         is_valid_index = is_valid_index_factory(LL)
972
973         idx = 0
974
975         # The first leaf MAY be a '+' symbol...
976         if is_valid_index(idx) and LL[idx].type == token.PLUS:
977             idx += 1
978
979         # The next/first leaf MAY be an empty LPAR...
980         if is_valid_index(idx) and is_empty_lpar(LL[idx]):
981             idx += 1
982
983         # The next/first leaf MUST be a string...
984         if not is_valid_index(idx) or LL[idx].type != token.STRING:
985             return TErr("Line does not start with a string.")
986
987         string_idx = idx
988
989         # Skip the string trailer, if one exists.
990         string_parser = StringParser()
991         idx = string_parser.parse(LL, string_idx)
992
993         # That string MAY be followed by an empty RPAR...
994         if is_valid_index(idx) and is_empty_rpar(LL[idx]):
995             idx += 1
996
997         # That string / empty RPAR leaf MAY be followed by a comma...
998         if is_valid_index(idx) and LL[idx].type == token.COMMA:
999             idx += 1
1000
1001         # But no more leaves are allowed...
1002         if is_valid_index(idx):
1003             return TErr("This line does not end with a string.")
1004
1005         return Ok(string_idx)
1006
1007     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
1008         LL = line.leaves
1009
1010         QUOTE = LL[string_idx].value[-1]
1011
1012         is_valid_index = is_valid_index_factory(LL)
1013         insert_str_child = insert_str_child_factory(LL[string_idx])
1014
1015         prefix = get_string_prefix(LL[string_idx].value)
1016
1017         # We MAY choose to drop the 'f' prefix from substrings that don't
1018         # contain any f-expressions, but ONLY if the original f-string
1019         # contains at least one f-expression. Otherwise, we will alter the AST
1020         # of the program.
1021         drop_pointless_f_prefix = ("f" in prefix) and re.search(
1022             self.RE_FEXPR, LL[string_idx].value, re.VERBOSE
1023         )
1024
1025         first_string_line = True
1026         starts_with_plus = LL[0].type == token.PLUS
1027
1028         def line_needs_plus() -> bool:
1029             return first_string_line and starts_with_plus
1030
1031         def maybe_append_plus(new_line: Line) -> None:
1032             """
1033             Side Effects:
1034                 If @line starts with a plus and this is the first line we are
1035                 constructing, this function appends a PLUS leaf to @new_line
1036                 and replaces the old PLUS leaf in the node structure. Otherwise
1037                 this function does nothing.
1038             """
1039             if line_needs_plus():
1040                 plus_leaf = Leaf(token.PLUS, "+")
1041                 replace_child(LL[0], plus_leaf)
1042                 new_line.append(plus_leaf)
1043
1044         ends_with_comma = (
1045             is_valid_index(string_idx + 1) and LL[string_idx + 1].type == token.COMMA
1046         )
1047
1048         def max_last_string() -> int:
1049             """
1050             Returns:
1051                 The max allowed length of the string value used for the last
1052                 line we will construct.
1053             """
1054             result = self.line_length
1055             result -= line.depth * 4
1056             result -= 1 if ends_with_comma else 0
1057             result -= 2 if line_needs_plus() else 0
1058             return result
1059
1060         # --- Calculate Max Break Index (for string value)
1061         # We start with the line length limit
1062         max_break_idx = self.line_length
1063         # The last index of a string of length N is N-1.
1064         max_break_idx -= 1
1065         # Leading whitespace is not present in the string value (e.g. Leaf.value).
1066         max_break_idx -= line.depth * 4
1067         if max_break_idx < 0:
1068             yield TErr(
1069                 f"Unable to split {LL[string_idx].value} at such high of a line depth:"
1070                 f" {line.depth}"
1071             )
1072             return
1073
1074         # Check if StringMerger registered any custom splits.
1075         custom_splits = self.pop_custom_splits(LL[string_idx].value)
1076         # We use them ONLY if none of them would produce lines that exceed the
1077         # line limit.
1078         use_custom_breakpoints = bool(
1079             custom_splits
1080             and all(csplit.break_idx <= max_break_idx for csplit in custom_splits)
1081         )
1082
1083         # Temporary storage for the remaining chunk of the string line that
1084         # can't fit onto the line currently being constructed.
1085         rest_value = LL[string_idx].value
1086
1087         def more_splits_should_be_made() -> bool:
1088             """
1089             Returns:
1090                 True iff `rest_value` (the remaining string value from the last
1091                 split), should be split again.
1092             """
1093             if use_custom_breakpoints:
1094                 return len(custom_splits) > 1
1095             else:
1096                 return len(rest_value) > max_last_string()
1097
1098         string_line_results: List[Ok[Line]] = []
1099         while more_splits_should_be_made():
1100             if use_custom_breakpoints:
1101                 # Custom User Split (manual)
1102                 csplit = custom_splits.pop(0)
1103                 break_idx = csplit.break_idx
1104             else:
1105                 # Algorithmic Split (automatic)
1106                 max_bidx = max_break_idx - 2 if line_needs_plus() else max_break_idx
1107                 maybe_break_idx = self._get_break_idx(rest_value, max_bidx)
1108                 if maybe_break_idx is None:
1109                     # If we are unable to algorithmically determine a good split
1110                     # and this string has custom splits registered to it, we
1111                     # fall back to using them--which means we have to start
1112                     # over from the beginning.
1113                     if custom_splits:
1114                         rest_value = LL[string_idx].value
1115                         string_line_results = []
1116                         first_string_line = True
1117                         use_custom_breakpoints = True
1118                         continue
1119
1120                     # Otherwise, we stop splitting here.
1121                     break
1122
1123                 break_idx = maybe_break_idx
1124
1125             # --- Construct `next_value`
1126             next_value = rest_value[:break_idx] + QUOTE
1127             if (
1128                 # Are we allowed to try to drop a pointless 'f' prefix?
1129                 drop_pointless_f_prefix
1130                 # If we are, will we be successful?
1131                 and next_value != self._normalize_f_string(next_value, prefix)
1132             ):
1133                 # If the current custom split did NOT originally use a prefix,
1134                 # then `csplit.break_idx` will be off by one after removing
1135                 # the 'f' prefix.
1136                 break_idx = (
1137                     break_idx + 1
1138                     if use_custom_breakpoints and not csplit.has_prefix
1139                     else break_idx
1140                 )
1141                 next_value = rest_value[:break_idx] + QUOTE
1142                 next_value = self._normalize_f_string(next_value, prefix)
1143
1144             # --- Construct `next_leaf`
1145             next_leaf = Leaf(token.STRING, next_value)
1146             insert_str_child(next_leaf)
1147             self._maybe_normalize_string_quotes(next_leaf)
1148
1149             # --- Construct `next_line`
1150             next_line = line.clone()
1151             maybe_append_plus(next_line)
1152             next_line.append(next_leaf)
1153             string_line_results.append(Ok(next_line))
1154
1155             rest_value = prefix + QUOTE + rest_value[break_idx:]
1156             first_string_line = False
1157
1158         yield from string_line_results
1159
1160         if drop_pointless_f_prefix:
1161             rest_value = self._normalize_f_string(rest_value, prefix)
1162
1163         rest_leaf = Leaf(token.STRING, rest_value)
1164         insert_str_child(rest_leaf)
1165
1166         # NOTE: I could not find a test case that verifies that the following
1167         # line is actually necessary, but it seems to be. Otherwise we risk
1168         # not normalizing the last substring, right?
1169         self._maybe_normalize_string_quotes(rest_leaf)
1170
1171         last_line = line.clone()
1172         maybe_append_plus(last_line)
1173
1174         # If there are any leaves to the right of the target string...
1175         if is_valid_index(string_idx + 1):
1176             # We use `temp_value` here to determine how long the last line
1177             # would be if we were to append all the leaves to the right of the
1178             # target string to the last string line.
1179             temp_value = rest_value
1180             for leaf in LL[string_idx + 1 :]:
1181                 temp_value += str(leaf)
1182                 if leaf.type == token.LPAR:
1183                     break
1184
1185             # Try to fit them all on the same line with the last substring...
1186             if (
1187                 len(temp_value) <= max_last_string()
1188                 or LL[string_idx + 1].type == token.COMMA
1189             ):
1190                 last_line.append(rest_leaf)
1191                 append_leaves(last_line, line, LL[string_idx + 1 :])
1192                 yield Ok(last_line)
1193             # Otherwise, place the last substring on one line and everything
1194             # else on a line below that...
1195             else:
1196                 last_line.append(rest_leaf)
1197                 yield Ok(last_line)
1198
1199                 non_string_line = line.clone()
1200                 append_leaves(non_string_line, line, LL[string_idx + 1 :])
1201                 yield Ok(non_string_line)
1202         # Else the target string was the last leaf...
1203         else:
1204             last_line.append(rest_leaf)
1205             last_line.comments = line.comments.copy()
1206             yield Ok(last_line)
1207
1208     def _get_break_idx(self, string: str, max_break_idx: int) -> Optional[int]:
1209         """
1210         This method contains the algorithm that StringSplitter uses to
1211         determine which character to split each string at.
1212
1213         Args:
1214             @string: The substring that we are attempting to split.
1215             @max_break_idx: The ideal break index. We will return this value if it
1216             meets all the necessary conditions. In the likely event that it
1217             doesn't we will try to find the closest index BELOW @max_break_idx
1218             that does. If that fails, we will expand our search by also
1219             considering all valid indices ABOVE @max_break_idx.
1220
1221         Pre-Conditions:
1222             * assert_is_leaf_string(@string)
1223             * 0 <= @max_break_idx < len(@string)
1224
1225         Returns:
1226             break_idx, if an index is able to be found that meets all of the
1227             conditions listed in the 'Transformations' section of this classes'
1228             docstring.
1229                 OR
1230             None, otherwise.
1231         """
1232         is_valid_index = is_valid_index_factory(string)
1233
1234         assert is_valid_index(max_break_idx)
1235         assert_is_leaf_string(string)
1236
1237         _fexpr_slices: Optional[List[Tuple[Index, Index]]] = None
1238
1239         def fexpr_slices() -> Iterator[Tuple[Index, Index]]:
1240             """
1241             Yields:
1242                 All ranges of @string which, if @string were to be split there,
1243                 would result in the splitting of an f-expression (which is NOT
1244                 allowed).
1245             """
1246             nonlocal _fexpr_slices
1247
1248             if _fexpr_slices is None:
1249                 _fexpr_slices = []
1250                 for match in re.finditer(self.RE_FEXPR, string, re.VERBOSE):
1251                     _fexpr_slices.append(match.span())
1252
1253             yield from _fexpr_slices
1254
1255         is_fstring = "f" in get_string_prefix(string)
1256
1257         def breaks_fstring_expression(i: Index) -> bool:
1258             """
1259             Returns:
1260                 True iff returning @i would result in the splitting of an
1261                 f-expression (which is NOT allowed).
1262             """
1263             if not is_fstring:
1264                 return False
1265
1266             for (start, end) in fexpr_slices():
1267                 if start <= i < end:
1268                     return True
1269
1270             return False
1271
1272         def passes_all_checks(i: Index) -> bool:
1273             """
1274             Returns:
1275                 True iff ALL of the conditions listed in the 'Transformations'
1276                 section of this classes' docstring would be be met by returning @i.
1277             """
1278             is_space = string[i] == " "
1279
1280             is_not_escaped = True
1281             j = i - 1
1282             while is_valid_index(j) and string[j] == "\\":
1283                 is_not_escaped = not is_not_escaped
1284                 j -= 1
1285
1286             is_big_enough = (
1287                 len(string[i:]) >= self.MIN_SUBSTR_SIZE
1288                 and len(string[:i]) >= self.MIN_SUBSTR_SIZE
1289             )
1290             return (
1291                 is_space
1292                 and is_not_escaped
1293                 and is_big_enough
1294                 and not breaks_fstring_expression(i)
1295             )
1296
1297         # First, we check all indices BELOW @max_break_idx.
1298         break_idx = max_break_idx
1299         while is_valid_index(break_idx - 1) and not passes_all_checks(break_idx):
1300             break_idx -= 1
1301
1302         if not passes_all_checks(break_idx):
1303             # If that fails, we check all indices ABOVE @max_break_idx.
1304             #
1305             # If we are able to find a valid index here, the next line is going
1306             # to be longer than the specified line length, but it's probably
1307             # better than doing nothing at all.
1308             break_idx = max_break_idx + 1
1309             while is_valid_index(break_idx + 1) and not passes_all_checks(break_idx):
1310                 break_idx += 1
1311
1312             if not is_valid_index(break_idx) or not passes_all_checks(break_idx):
1313                 return None
1314
1315         return break_idx
1316
1317     def _maybe_normalize_string_quotes(self, leaf: Leaf) -> None:
1318         if self.normalize_strings:
1319             leaf.value = normalize_string_quotes(leaf.value)
1320
1321     def _normalize_f_string(self, string: str, prefix: str) -> str:
1322         """
1323         Pre-Conditions:
1324             * assert_is_leaf_string(@string)
1325
1326         Returns:
1327             * If @string is an f-string that contains no f-expressions, we
1328             return a string identical to @string except that the 'f' prefix
1329             has been stripped and all double braces (i.e. '{{' or '}}') have
1330             been normalized (i.e. turned into '{' or '}').
1331                 OR
1332             * Otherwise, we return @string.
1333         """
1334         assert_is_leaf_string(string)
1335
1336         if "f" in prefix and not re.search(self.RE_FEXPR, string, re.VERBOSE):
1337             new_prefix = prefix.replace("f", "")
1338
1339             temp = string[len(prefix) :]
1340             temp = re.sub(r"\{\{", "{", temp)
1341             temp = re.sub(r"\}\}", "}", temp)
1342             new_string = temp
1343
1344             return f"{new_prefix}{new_string}"
1345         else:
1346             return string
1347
1348
1349 class StringParenWrapper(CustomSplitMapMixin, BaseStringSplitter):
1350     """
1351     StringTransformer that splits non-"atom" strings (i.e. strings that do not
1352     exist on lines by themselves).
1353
1354     Requirements:
1355         All of the requirements listed in BaseStringSplitter's docstring in
1356         addition to the requirements listed below:
1357
1358         * The line is a return/yield statement, which returns/yields a string.
1359             OR
1360         * The line is part of a ternary expression (e.g. `x = y if cond else
1361         z`) such that the line starts with `else <string>`, where <string> is
1362         some string.
1363             OR
1364         * The line is an assert statement, which ends with a string.
1365             OR
1366         * The line is an assignment statement (e.g. `x = <string>` or `x +=
1367         <string>`) such that the variable is being assigned the value of some
1368         string.
1369             OR
1370         * The line is a dictionary key assignment where some valid key is being
1371         assigned the value of some string.
1372
1373     Transformations:
1374         The chosen string is wrapped in parentheses and then split at the LPAR.
1375
1376         We then have one line which ends with an LPAR and another line that
1377         starts with the chosen string. The latter line is then split again at
1378         the RPAR. This results in the RPAR (and possibly a trailing comma)
1379         being placed on its own line.
1380
1381         NOTE: If any leaves exist to the right of the chosen string (except
1382         for a trailing comma, which would be placed after the RPAR), those
1383         leaves are placed inside the parentheses.  In effect, the chosen
1384         string is not necessarily being "wrapped" by parentheses. We can,
1385         however, count on the LPAR being placed directly before the chosen
1386         string.
1387
1388         In other words, StringParenWrapper creates "atom" strings. These
1389         can then be split again by StringSplitter, if necessary.
1390
1391     Collaborations:
1392         In the event that a string line split by StringParenWrapper is
1393         changed such that it no longer needs to be given its own line,
1394         StringParenWrapper relies on StringParenStripper to clean up the
1395         parentheses it created.
1396     """
1397
1398     def do_splitter_match(self, line: Line) -> TMatchResult:
1399         LL = line.leaves
1400
1401         string_idx = (
1402             self._return_match(LL)
1403             or self._else_match(LL)
1404             or self._assert_match(LL)
1405             or self._assign_match(LL)
1406             or self._dict_match(LL)
1407         )
1408
1409         if string_idx is not None:
1410             string_value = line.leaves[string_idx].value
1411             # If the string has no spaces...
1412             if " " not in string_value:
1413                 # And will still violate the line length limit when split...
1414                 max_string_length = self.line_length - ((line.depth + 1) * 4)
1415                 if len(string_value) > max_string_length:
1416                     # And has no associated custom splits...
1417                     if not self.has_custom_splits(string_value):
1418                         # Then we should NOT put this string on its own line.
1419                         return TErr(
1420                             "We do not wrap long strings in parentheses when the"
1421                             " resultant line would still be over the specified line"
1422                             " length and can't be split further by StringSplitter."
1423                         )
1424             return Ok(string_idx)
1425
1426         return TErr("This line does not contain any non-atomic strings.")
1427
1428     @staticmethod
1429     def _return_match(LL: List[Leaf]) -> Optional[int]:
1430         """
1431         Returns:
1432             string_idx such that @LL[string_idx] is equal to our target (i.e.
1433             matched) string, if this line matches the return/yield statement
1434             requirements listed in the 'Requirements' section of this classes'
1435             docstring.
1436                 OR
1437             None, otherwise.
1438         """
1439         # If this line is apart of a return/yield statement and the first leaf
1440         # contains either the "return" or "yield" keywords...
1441         if parent_type(LL[0]) in [syms.return_stmt, syms.yield_expr] and LL[
1442             0
1443         ].value in ["return", "yield"]:
1444             is_valid_index = is_valid_index_factory(LL)
1445
1446             idx = 2 if is_valid_index(1) and is_empty_par(LL[1]) else 1
1447             # The next visible leaf MUST contain a string...
1448             if is_valid_index(idx) and LL[idx].type == token.STRING:
1449                 return idx
1450
1451         return None
1452
1453     @staticmethod
1454     def _else_match(LL: List[Leaf]) -> Optional[int]:
1455         """
1456         Returns:
1457             string_idx such that @LL[string_idx] is equal to our target (i.e.
1458             matched) string, if this line matches the ternary expression
1459             requirements listed in the 'Requirements' section of this classes'
1460             docstring.
1461                 OR
1462             None, otherwise.
1463         """
1464         # If this line is apart of a ternary expression and the first leaf
1465         # contains the "else" keyword...
1466         if (
1467             parent_type(LL[0]) == syms.test
1468             and LL[0].type == token.NAME
1469             and LL[0].value == "else"
1470         ):
1471             is_valid_index = is_valid_index_factory(LL)
1472
1473             idx = 2 if is_valid_index(1) and is_empty_par(LL[1]) else 1
1474             # The next visible leaf MUST contain a string...
1475             if is_valid_index(idx) and LL[idx].type == token.STRING:
1476                 return idx
1477
1478         return None
1479
1480     @staticmethod
1481     def _assert_match(LL: List[Leaf]) -> Optional[int]:
1482         """
1483         Returns:
1484             string_idx such that @LL[string_idx] is equal to our target (i.e.
1485             matched) string, if this line matches the assert statement
1486             requirements listed in the 'Requirements' section of this classes'
1487             docstring.
1488                 OR
1489             None, otherwise.
1490         """
1491         # If this line is apart of an assert statement and the first leaf
1492         # contains the "assert" keyword...
1493         if parent_type(LL[0]) == syms.assert_stmt and LL[0].value == "assert":
1494             is_valid_index = is_valid_index_factory(LL)
1495
1496             for (i, leaf) in enumerate(LL):
1497                 # We MUST find a comma...
1498                 if leaf.type == token.COMMA:
1499                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
1500
1501                     # That comma MUST be followed by a string...
1502                     if is_valid_index(idx) and LL[idx].type == token.STRING:
1503                         string_idx = idx
1504
1505                         # Skip the string trailer, if one exists.
1506                         string_parser = StringParser()
1507                         idx = string_parser.parse(LL, string_idx)
1508
1509                         # But no more leaves are allowed...
1510                         if not is_valid_index(idx):
1511                             return string_idx
1512
1513         return None
1514
1515     @staticmethod
1516     def _assign_match(LL: List[Leaf]) -> Optional[int]:
1517         """
1518         Returns:
1519             string_idx such that @LL[string_idx] is equal to our target (i.e.
1520             matched) string, if this line matches the assignment statement
1521             requirements listed in the 'Requirements' section of this classes'
1522             docstring.
1523                 OR
1524             None, otherwise.
1525         """
1526         # If this line is apart of an expression statement or is a function
1527         # argument AND the first leaf contains a variable name...
1528         if (
1529             parent_type(LL[0]) in [syms.expr_stmt, syms.argument, syms.power]
1530             and LL[0].type == token.NAME
1531         ):
1532             is_valid_index = is_valid_index_factory(LL)
1533
1534             for (i, leaf) in enumerate(LL):
1535                 # We MUST find either an '=' or '+=' symbol...
1536                 if leaf.type in [token.EQUAL, token.PLUSEQUAL]:
1537                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
1538
1539                     # That symbol MUST be followed by a string...
1540                     if is_valid_index(idx) and LL[idx].type == token.STRING:
1541                         string_idx = idx
1542
1543                         # Skip the string trailer, if one exists.
1544                         string_parser = StringParser()
1545                         idx = string_parser.parse(LL, string_idx)
1546
1547                         # The next leaf MAY be a comma iff this line is apart
1548                         # of a function argument...
1549                         if (
1550                             parent_type(LL[0]) == syms.argument
1551                             and is_valid_index(idx)
1552                             and LL[idx].type == token.COMMA
1553                         ):
1554                             idx += 1
1555
1556                         # But no more leaves are allowed...
1557                         if not is_valid_index(idx):
1558                             return string_idx
1559
1560         return None
1561
1562     @staticmethod
1563     def _dict_match(LL: List[Leaf]) -> Optional[int]:
1564         """
1565         Returns:
1566             string_idx such that @LL[string_idx] is equal to our target (i.e.
1567             matched) string, if this line matches the dictionary key assignment
1568             statement requirements listed in the 'Requirements' section of this
1569             classes' docstring.
1570                 OR
1571             None, otherwise.
1572         """
1573         # If this line is apart of a dictionary key assignment...
1574         if syms.dictsetmaker in [parent_type(LL[0]), parent_type(LL[0].parent)]:
1575             is_valid_index = is_valid_index_factory(LL)
1576
1577             for (i, leaf) in enumerate(LL):
1578                 # We MUST find a colon...
1579                 if leaf.type == token.COLON:
1580                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
1581
1582                     # That colon MUST be followed by a string...
1583                     if is_valid_index(idx) and LL[idx].type == token.STRING:
1584                         string_idx = idx
1585
1586                         # Skip the string trailer, if one exists.
1587                         string_parser = StringParser()
1588                         idx = string_parser.parse(LL, string_idx)
1589
1590                         # That string MAY be followed by a comma...
1591                         if is_valid_index(idx) and LL[idx].type == token.COMMA:
1592                             idx += 1
1593
1594                         # But no more leaves are allowed...
1595                         if not is_valid_index(idx):
1596                             return string_idx
1597
1598         return None
1599
1600     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
1601         LL = line.leaves
1602
1603         is_valid_index = is_valid_index_factory(LL)
1604         insert_str_child = insert_str_child_factory(LL[string_idx])
1605
1606         comma_idx = -1
1607         ends_with_comma = False
1608         if LL[comma_idx].type == token.COMMA:
1609             ends_with_comma = True
1610
1611         leaves_to_steal_comments_from = [LL[string_idx]]
1612         if ends_with_comma:
1613             leaves_to_steal_comments_from.append(LL[comma_idx])
1614
1615         # --- First Line
1616         first_line = line.clone()
1617         left_leaves = LL[:string_idx]
1618
1619         # We have to remember to account for (possibly invisible) LPAR and RPAR
1620         # leaves that already wrapped the target string. If these leaves do
1621         # exist, we will replace them with our own LPAR and RPAR leaves.
1622         old_parens_exist = False
1623         if left_leaves and left_leaves[-1].type == token.LPAR:
1624             old_parens_exist = True
1625             leaves_to_steal_comments_from.append(left_leaves[-1])
1626             left_leaves.pop()
1627
1628         append_leaves(first_line, line, left_leaves)
1629
1630         lpar_leaf = Leaf(token.LPAR, "(")
1631         if old_parens_exist:
1632             replace_child(LL[string_idx - 1], lpar_leaf)
1633         else:
1634             insert_str_child(lpar_leaf)
1635         first_line.append(lpar_leaf)
1636
1637         # We throw inline comments that were originally to the right of the
1638         # target string to the top line. They will now be shown to the right of
1639         # the LPAR.
1640         for leaf in leaves_to_steal_comments_from:
1641             for comment_leaf in line.comments_after(leaf):
1642                 first_line.append(comment_leaf, preformatted=True)
1643
1644         yield Ok(first_line)
1645
1646         # --- Middle (String) Line
1647         # We only need to yield one (possibly too long) string line, since the
1648         # `StringSplitter` will break it down further if necessary.
1649         string_value = LL[string_idx].value
1650         string_line = Line(
1651             mode=line.mode,
1652             depth=line.depth + 1,
1653             inside_brackets=True,
1654             should_split_rhs=line.should_split_rhs,
1655             magic_trailing_comma=line.magic_trailing_comma,
1656         )
1657         string_leaf = Leaf(token.STRING, string_value)
1658         insert_str_child(string_leaf)
1659         string_line.append(string_leaf)
1660
1661         old_rpar_leaf = None
1662         if is_valid_index(string_idx + 1):
1663             right_leaves = LL[string_idx + 1 :]
1664             if ends_with_comma:
1665                 right_leaves.pop()
1666
1667             if old_parens_exist:
1668                 assert (
1669                     right_leaves and right_leaves[-1].type == token.RPAR
1670                 ), "Apparently, old parentheses do NOT exist?!"
1671                 old_rpar_leaf = right_leaves.pop()
1672
1673             append_leaves(string_line, line, right_leaves)
1674
1675         yield Ok(string_line)
1676
1677         # --- Last Line
1678         last_line = line.clone()
1679         last_line.bracket_tracker = first_line.bracket_tracker
1680
1681         new_rpar_leaf = Leaf(token.RPAR, ")")
1682         if old_rpar_leaf is not None:
1683             replace_child(old_rpar_leaf, new_rpar_leaf)
1684         else:
1685             insert_str_child(new_rpar_leaf)
1686         last_line.append(new_rpar_leaf)
1687
1688         # If the target string ended with a comma, we place this comma to the
1689         # right of the RPAR on the last line.
1690         if ends_with_comma:
1691             comma_leaf = Leaf(token.COMMA, ",")
1692             replace_child(LL[comma_idx], comma_leaf)
1693             last_line.append(comma_leaf)
1694
1695         yield Ok(last_line)
1696
1697
1698 class StringParser:
1699     """
1700     A state machine that aids in parsing a string's "trailer", which can be
1701     either non-existent, an old-style formatting sequence (e.g. `% varX` or `%
1702     (varX, varY)`), or a method-call / attribute access (e.g. `.format(varX,
1703     varY)`).
1704
1705     NOTE: A new StringParser object MUST be instantiated for each string
1706     trailer we need to parse.
1707
1708     Examples:
1709         We shall assume that `line` equals the `Line` object that corresponds
1710         to the following line of python code:
1711         ```
1712         x = "Some {}.".format("String") + some_other_string
1713         ```
1714
1715         Furthermore, we will assume that `string_idx` is some index such that:
1716         ```
1717         assert line.leaves[string_idx].value == "Some {}."
1718         ```
1719
1720         The following code snippet then holds:
1721         ```
1722         string_parser = StringParser()
1723         idx = string_parser.parse(line.leaves, string_idx)
1724         assert line.leaves[idx].type == token.PLUS
1725         ```
1726     """
1727
1728     DEFAULT_TOKEN = -1
1729
1730     # String Parser States
1731     START = 1
1732     DOT = 2
1733     NAME = 3
1734     PERCENT = 4
1735     SINGLE_FMT_ARG = 5
1736     LPAR = 6
1737     RPAR = 7
1738     DONE = 8
1739
1740     # Lookup Table for Next State
1741     _goto: Dict[Tuple[ParserState, NodeType], ParserState] = {
1742         # A string trailer may start with '.' OR '%'.
1743         (START, token.DOT): DOT,
1744         (START, token.PERCENT): PERCENT,
1745         (START, DEFAULT_TOKEN): DONE,
1746         # A '.' MUST be followed by an attribute or method name.
1747         (DOT, token.NAME): NAME,
1748         # A method name MUST be followed by an '(', whereas an attribute name
1749         # is the last symbol in the string trailer.
1750         (NAME, token.LPAR): LPAR,
1751         (NAME, DEFAULT_TOKEN): DONE,
1752         # A '%' symbol can be followed by an '(' or a single argument (e.g. a
1753         # string or variable name).
1754         (PERCENT, token.LPAR): LPAR,
1755         (PERCENT, DEFAULT_TOKEN): SINGLE_FMT_ARG,
1756         # If a '%' symbol is followed by a single argument, that argument is
1757         # the last leaf in the string trailer.
1758         (SINGLE_FMT_ARG, DEFAULT_TOKEN): DONE,
1759         # If present, a ')' symbol is the last symbol in a string trailer.
1760         # (NOTE: LPARS and nested RPARS are not included in this lookup table,
1761         # since they are treated as a special case by the parsing logic in this
1762         # classes' implementation.)
1763         (RPAR, DEFAULT_TOKEN): DONE,
1764     }
1765
1766     def __init__(self) -> None:
1767         self._state = self.START
1768         self._unmatched_lpars = 0
1769
1770     def parse(self, leaves: List[Leaf], string_idx: int) -> int:
1771         """
1772         Pre-conditions:
1773             * @leaves[@string_idx].type == token.STRING
1774
1775         Returns:
1776             The index directly after the last leaf which is apart of the string
1777             trailer, if a "trailer" exists.
1778                 OR
1779             @string_idx + 1, if no string "trailer" exists.
1780         """
1781         assert leaves[string_idx].type == token.STRING
1782
1783         idx = string_idx + 1
1784         while idx < len(leaves) and self._next_state(leaves[idx]):
1785             idx += 1
1786         return idx
1787
1788     def _next_state(self, leaf: Leaf) -> bool:
1789         """
1790         Pre-conditions:
1791             * On the first call to this function, @leaf MUST be the leaf that
1792             was directly after the string leaf in question (e.g. if our target
1793             string is `line.leaves[i]` then the first call to this method must
1794             be `line.leaves[i + 1]`).
1795             * On the next call to this function, the leaf parameter passed in
1796             MUST be the leaf directly following @leaf.
1797
1798         Returns:
1799             True iff @leaf is apart of the string's trailer.
1800         """
1801         # We ignore empty LPAR or RPAR leaves.
1802         if is_empty_par(leaf):
1803             return True
1804
1805         next_token = leaf.type
1806         if next_token == token.LPAR:
1807             self._unmatched_lpars += 1
1808
1809         current_state = self._state
1810
1811         # The LPAR parser state is a special case. We will return True until we
1812         # find the matching RPAR token.
1813         if current_state == self.LPAR:
1814             if next_token == token.RPAR:
1815                 self._unmatched_lpars -= 1
1816                 if self._unmatched_lpars == 0:
1817                     self._state = self.RPAR
1818         # Otherwise, we use a lookup table to determine the next state.
1819         else:
1820             # If the lookup table matches the current state to the next
1821             # token, we use the lookup table.
1822             if (current_state, next_token) in self._goto:
1823                 self._state = self._goto[current_state, next_token]
1824             else:
1825                 # Otherwise, we check if a the current state was assigned a
1826                 # default.
1827                 if (current_state, self.DEFAULT_TOKEN) in self._goto:
1828                     self._state = self._goto[current_state, self.DEFAULT_TOKEN]
1829                 # If no default has been assigned, then this parser has a logic
1830                 # error.
1831                 else:
1832                     raise RuntimeError(f"{self.__class__.__name__} LOGIC ERROR!")
1833
1834             if self._state == self.DONE:
1835                 return False
1836
1837         return True
1838
1839
1840 def insert_str_child_factory(string_leaf: Leaf) -> Callable[[LN], None]:
1841     """
1842     Factory for a convenience function that is used to orphan @string_leaf
1843     and then insert multiple new leaves into the same part of the node
1844     structure that @string_leaf had originally occupied.
1845
1846     Examples:
1847         Let `string_leaf = Leaf(token.STRING, '"foo"')` and `N =
1848         string_leaf.parent`. Assume the node `N` has the following
1849         original structure:
1850
1851         Node(
1852             expr_stmt, [
1853                 Leaf(NAME, 'x'),
1854                 Leaf(EQUAL, '='),
1855                 Leaf(STRING, '"foo"'),
1856             ]
1857         )
1858
1859         We then run the code snippet shown below.
1860         ```
1861         insert_str_child = insert_str_child_factory(string_leaf)
1862
1863         lpar = Leaf(token.LPAR, '(')
1864         insert_str_child(lpar)
1865
1866         bar = Leaf(token.STRING, '"bar"')
1867         insert_str_child(bar)
1868
1869         rpar = Leaf(token.RPAR, ')')
1870         insert_str_child(rpar)
1871         ```
1872
1873         After which point, it follows that `string_leaf.parent is None` and
1874         the node `N` now has the following structure:
1875
1876         Node(
1877             expr_stmt, [
1878                 Leaf(NAME, 'x'),
1879                 Leaf(EQUAL, '='),
1880                 Leaf(LPAR, '('),
1881                 Leaf(STRING, '"bar"'),
1882                 Leaf(RPAR, ')'),
1883             ]
1884         )
1885     """
1886     string_parent = string_leaf.parent
1887     string_child_idx = string_leaf.remove()
1888
1889     def insert_str_child(child: LN) -> None:
1890         nonlocal string_child_idx
1891
1892         assert string_parent is not None
1893         assert string_child_idx is not None
1894
1895         string_parent.insert_child(string_child_idx, child)
1896         string_child_idx += 1
1897
1898     return insert_str_child
1899
1900
1901 def is_valid_index_factory(seq: Sequence[Any]) -> Callable[[int], bool]:
1902     """
1903     Examples:
1904         ```
1905         my_list = [1, 2, 3]
1906
1907         is_valid_index = is_valid_index_factory(my_list)
1908
1909         assert is_valid_index(0)
1910         assert is_valid_index(2)
1911
1912         assert not is_valid_index(3)
1913         assert not is_valid_index(-1)
1914         ```
1915     """
1916
1917     def is_valid_index(idx: int) -> bool:
1918         """
1919         Returns:
1920             True iff @idx is positive AND seq[@idx] does NOT raise an
1921             IndexError.
1922         """
1923         return 0 <= idx < len(seq)
1924
1925     return is_valid_index