]> 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:

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