+def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[LeafID]]:
+ """Generate sets of closing bracket IDs that should be omitted in a RHS.
+
+ Brackets can be omitted if the entire trailer up to and including
+ a preceding closing bracket fits in one line.
+
+ Yielded sets are cumulative (contain results of previous yields, too). First
+ set is empty.
+ """
+
+ omit: Set[LeafID] = set()
+ yield omit
+
+ length = 4 * line.depth
+ opening_bracket = None
+ closing_bracket = None
+ optional_brackets: Set[LeafID] = set()
+ inner_brackets: Set[LeafID] = set()
+ for index, leaf in enumerate_reversed(line.leaves):
+ length += len(leaf.prefix) + len(leaf.value)
+ if length > line_length:
+ break
+
+ comment: Optional[Leaf]
+ for comment in line.comments_after(leaf, index):
+ if "\n" in comment.prefix:
+ break # Oops, standalone comment!
+
+ length += len(comment.value)
+ else:
+ comment = None
+ if comment is not None:
+ break # There was a standalone comment, we can't continue.
+
+ optional_brackets.discard(id(leaf))
+ if opening_bracket:
+ if leaf is opening_bracket:
+ opening_bracket = None
+ elif leaf.type in CLOSING_BRACKETS:
+ inner_brackets.add(id(leaf))
+ elif leaf.type in CLOSING_BRACKETS:
+ if not leaf.value:
+ optional_brackets.add(id(opening_bracket))
+ continue
+
+ if index > 0 and line.leaves[index - 1].type in OPENING_BRACKETS:
+ # Empty brackets would fail a split so treat them as "inner"
+ # brackets (e.g. only add them to the `omit` set if another
+ # pair of brackets was good enough.
+ inner_brackets.add(id(leaf))
+ continue
+
+ opening_bracket = leaf.opening_bracket
+ if closing_bracket:
+ omit.add(id(closing_bracket))
+ omit.update(inner_brackets)
+ inner_brackets.clear()
+ yield omit
+ closing_bracket = leaf
+
+