before = 0 if depth else 1
else:
before = 1 if depth else 2
- is_decorator = current_line.is_decorator
- if is_decorator or current_line.is_def or current_line.is_class:
- if not is_decorator:
- self.previous_defs.append(depth)
- if self.previous_line is None:
- # Don't insert empty lines before the first line in the file.
- return 0, 0
-
- if self.previous_line.is_decorator:
- return 0, 0
-
- if self.previous_line.depth < current_line.depth and (
- self.previous_line.is_class or self.previous_line.is_def
- ):
- return 0, 0
-
- if (
- self.previous_line.is_comment
- and self.previous_line.depth == current_line.depth
- and before == 0
- ):
- return 0, 0
-
- if self.is_pyi:
- if self.previous_line.depth > current_line.depth:
- newlines = 1
- elif current_line.is_class or self.previous_line.is_class:
- if current_line.is_stub_class and self.previous_line.is_stub_class:
- newlines = 0
- else:
- newlines = 1
- else:
- newlines = 0
- else:
- newlines = 2
- if current_line.depth and newlines:
- newlines -= 1
- return newlines, 0
+ if current_line.is_decorator or current_line.is_def or current_line.is_class:
+ return self._maybe_empty_lines_for_class_or_def(current_line, before)
if (
self.previous_line
return before, 0
+ def _maybe_empty_lines_for_class_or_def(
+ self, current_line: Line, before: int
+ ) -> Tuple[int, int]:
+ if not current_line.is_decorator:
+ self.previous_defs.append(current_line.depth)
+ if self.previous_line is None:
+ # Don't insert empty lines before the first line in the file.
+ return 0, 0
+
+ if self.previous_line.is_decorator:
+ return 0, 0
+
+ if self.previous_line.depth < current_line.depth and (
+ self.previous_line.is_class or self.previous_line.is_def
+ ):
+ return 0, 0
+
+ if (
+ self.previous_line.is_comment
+ and self.previous_line.depth == current_line.depth
+ and before == 0
+ ):
+ return 0, 0
+
+ if self.is_pyi:
+ if self.previous_line.depth > current_line.depth:
+ newlines = 1
+ elif current_line.is_class or self.previous_line.is_class:
+ if current_line.is_stub_class and self.previous_line.is_stub_class:
+ # No blank line between classes with an emty body
+ newlines = 0
+ else:
+ newlines = 1
+ elif current_line.is_def and not self.previous_line.is_def:
+ # Blank line between a block of functions and a block of non-functions
+ newlines = 1
+ else:
+ newlines = 0
+ else:
+ newlines = 2
+ if current_line.depth and newlines:
+ newlines -= 1
+ return newlines, 0
+
@dataclass
class LineGenerator(Visitor[Line]):
return 0
+FMT_OFF = {"# fmt: off", "# fmt:off", "# yapf: disable"}
+FMT_ON = {"# fmt: on", "# fmt:on", "# yapf: enable"}
+
+
def generate_comments(leaf: LN) -> Iterator[Leaf]:
"""Clean the prefix of the `leaf` and generate comments from it, if any.
Inline comments are emitted as regular token.COMMENT leaves. Standalone
are emitted with a fake STANDALONE_COMMENT token identifier.
"""
- p = leaf.prefix
- if not p:
- return
+ for pc in list_comments(leaf.prefix, is_endmarker=leaf.type == token.ENDMARKER):
+ yield Leaf(pc.type, pc.value, prefix="\n" * pc.newlines)
+ if pc.value in FMT_ON:
+ raise FormatOn(pc.consumed)
+
+ if pc.value in FMT_OFF:
+ if pc.type == STANDALONE_COMMENT:
+ raise FormatOff(pc.consumed)
+
+ prev = preceding_leaf(leaf)
+ if not prev or prev.type in WHITESPACE: # standalone comment in disguise
+ raise FormatOff(pc.consumed)
+
+
+@dataclass
+class ProtoComment:
+ type: int # token.COMMENT or STANDALONE_COMMENT
+ value: str # content of the comment
+ newlines: int # how many newlines before the comment
+ consumed: int # how many characters of the original leaf's prefix did we consume
- if "#" not in p:
- return
+
+@lru_cache(maxsize=4096)
+def list_comments(prefix: str, is_endmarker: bool) -> List[ProtoComment]:
+ result: List[ProtoComment] = []
+ if not prefix or "#" not in prefix:
+ return result
consumed = 0
nlines = 0
- for index, line in enumerate(p.split("\n")):
+ for index, line in enumerate(prefix.split("\n")):
consumed += len(line) + 1 # adding the length of the split '\n'
line = line.lstrip()
if not line:
if not line.startswith("#"):
continue
- if index == 0 and leaf.type != token.ENDMARKER:
+ if index == 0 and not is_endmarker:
comment_type = token.COMMENT # simple trailing comment
else:
comment_type = STANDALONE_COMMENT
comment = make_comment(line)
- yield Leaf(comment_type, comment, prefix="\n" * nlines)
-
- if comment in {"# fmt: on", "# yapf: enable"}:
- raise FormatOn(consumed)
-
- if comment in {"# fmt: off", "# yapf: disable"}:
- if comment_type == STANDALONE_COMMENT:
- raise FormatOff(consumed)
-
- prev = preceding_leaf(leaf)
- if not prev or prev.type in WHITESPACE: # standalone comment in disguise
- raise FormatOff(consumed)
-
+ result.append(
+ ProtoComment(
+ type=comment_type, value=comment, newlines=nlines, consumed=consumed
+ )
+ )
nlines = 0
+ return result
def make_comment(content: str) -> str:
Standardizes on visible parentheses for single-element tuples, and keeps
existing visible parentheses for other tuples and generator expressions.
"""
- try:
- list(generate_comments(node))
- except FormatOff:
- return # This `node` has a prefix with `# fmt: off`, don't mess with parens.
+ for pc in list_comments(node.prefix, is_endmarker=False):
+ if pc.value in FMT_OFF:
+ # This `node` has a prefix with `# fmt: off`, don't mess with parens.
+ return
check_lpar = False
for index, child in enumerate(list(node.children)):