+
+def is_stub_suite(node: Node) -> bool:
+ """Return True if `node` is a suite with a stub body."""
+ if (
+ len(node.children) != 4
+ or node.children[0].type != token.NEWLINE
+ or node.children[1].type != token.INDENT
+ or node.children[3].type != token.DEDENT
+ ):
+ return False
+
+ return is_stub_body(node.children[2])
+
+
+def is_stub_body(node: LN) -> bool:
+ """Return True if `node` is a simple statement containing an ellipsis."""
+ if not isinstance(node, Node) or node.type != syms.simple_stmt:
+ return False
+
+ if len(node.children) != 2:
+ return False
+
+ child = node.children[0]
+ return (
+ child.type == syms.atom
+ and len(child.children) == 3
+ and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
+ )
+
+
+def max_delimiter_priority_in_atom(node: LN) -> int:
+ """Return maximum delimiter priority inside `node`.
+
+ This is specific to atoms with contents contained in a pair of parentheses.
+ If `node` isn't an atom or there are no enclosing parentheses, returns 0.
+ """
+ if node.type != syms.atom:
+ return 0
+
+ first = node.children[0]
+ last = node.children[-1]
+ if not (first.type == token.LPAR and last.type == token.RPAR):
+ return 0
+
+ bt = BracketTracker()
+ for c in node.children[1:-1]:
+ if isinstance(c, Leaf):
+ bt.mark(c)
+ else:
+ for leaf in c.leaves():
+ bt.mark(leaf)
+ try:
+ return bt.max_delimiter_priority()
+
+ except ValueError:
+ return 0
+
+
+def ensure_visible(leaf: Leaf) -> None:
+ """Make sure parentheses are visible.
+
+ They could be invisible as part of some statements (see
+ :func:`normalize_invible_parens` and :func:`visit_import_from`).
+ """
+ if leaf.type == token.LPAR:
+ leaf.value = "("
+ elif leaf.type == token.RPAR:
+ leaf.value = ")"
+
+
+def should_explode(line: Line, opening_bracket: Leaf) -> bool:
+ """Should `line` immediately be split with `delimiter_split()` after RHS?"""
+
+ if not (
+ opening_bracket.parent
+ and opening_bracket.parent.type in {syms.atom, syms.import_from}
+ and opening_bracket.value in "[{("
+ ):
+ return False
+
+ try:
+ last_leaf = line.leaves[-1]
+ exclude = {id(last_leaf)} if last_leaf.type == token.COMMA else set()
+ max_priority = line.bracket_tracker.max_delimiter_priority(exclude=exclude)
+ except (IndexError, ValueError):
+ return False
+
+ return max_priority == COMMA_PRIORITY
+
+
+def is_python36(node: Node) -> bool:
+ """Return True if the current file is using Python 3.6+ features.
+
+ Currently looking for:
+ - f-strings;
+ - underscores in numeric literals; and
+ - trailing commas after * or ** in function signatures and calls.
+ """
+ for n in node.pre_order():
+ if n.type == token.STRING:
+ value_head = n.value[:2] # type: ignore
+ if value_head in {'f"', 'F"', "f'", "F'", "rf", "fr", "RF", "FR"}:
+ return True
+
+ elif n.type == token.NUMBER:
+ if "_" in n.value: # type: ignore
+ return True
+
+ elif (
+ n.type in {syms.typedargslist, syms.arglist}
+ and n.children
+ and n.children[-1].type == token.COMMA
+ ):
+ for ch in n.children:
+ if ch.type in STARS:
+ return True
+
+ if ch.type == syms.argument:
+ for argch in ch.children:
+ if argch.type in STARS:
+ return True
+
+ return False
+
+
+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
+ inner_brackets: Set[LeafID] = set()
+ for index, leaf, leaf_length in enumerate_with_length(line, reversed=True):
+ length += leaf_length
+ if length > line_length:
+ break
+
+ has_inline_comment = leaf_length > len(leaf.value) + len(leaf.prefix)
+ if leaf.type == STANDALONE_COMMENT or has_inline_comment:
+ break
+
+ 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 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
+
+ if closing_bracket:
+ omit.add(id(closing_bracket))
+ omit.update(inner_brackets)
+ inner_brackets.clear()
+ yield omit
+
+ if leaf.value:
+ opening_bracket = leaf.opening_bracket
+ closing_bracket = leaf
+
+
+def get_future_imports(node: Node) -> Set[str]:
+ """Return a set of __future__ imports in the file."""
+ imports: Set[str] = set()
+
+ def get_imports_from_children(children: List[LN]) -> Generator[str, None, None]:
+ for child in children:
+ if isinstance(child, Leaf):
+ if child.type == token.NAME:
+ yield child.value
+ elif child.type == syms.import_as_name:
+ orig_name = child.children[0]
+ assert isinstance(orig_name, Leaf), "Invalid syntax parsing imports"
+ assert orig_name.type == token.NAME, "Invalid syntax parsing imports"
+ yield orig_name.value
+ elif child.type == syms.import_as_names:
+ yield from get_imports_from_children(child.children)
+ else:
+ assert False, "Invalid syntax parsing imports"
+
+ for child in node.children:
+ if child.type != syms.simple_stmt:
+ break
+ first_child = child.children[0]
+ if isinstance(first_child, Leaf):
+ # Continue looking if we see a docstring; otherwise stop.
+ if (
+ len(child.children) == 2
+ and first_child.type == token.STRING
+ and child.children[1].type == token.NEWLINE
+ ):
+ continue
+ else:
+ break
+ elif first_child.type == syms.import_from:
+ module_name = first_child.children[1]
+ if not isinstance(module_name, Leaf) or module_name.value != "__future__":
+ break
+ imports |= set(get_imports_from_children(first_child.children[3:]))
+ else:
+ break
+ return imports
+
+
+def gen_python_files_in_dir(
+ path: Path,
+ root: Path,
+ include: Pattern[str],
+ exclude: Pattern[str],
+ report: "Report",
+) -> Iterator[Path]:
+ """Generate all files under `path` whose paths are not excluded by the
+ `exclude` regex, but are included by the `include` regex.
+
+ Symbolic links pointing outside of the `root` directory are ignored.
+
+ `report` is where output about exclusions goes.
+ """
+ assert root.is_absolute(), f"INTERNAL ERROR: `root` must be absolute but is {root}"
+ for child in path.iterdir():
+ try:
+ normalized_path = "/" + child.resolve().relative_to(root).as_posix()
+ except ValueError:
+ if child.is_symlink():
+ report.path_ignored(
+ child, f"is a symbolic link that points outside {root}"
+ )
+ continue
+
+ raise
+
+ if child.is_dir():
+ normalized_path += "/"
+ exclude_match = exclude.search(normalized_path)
+ if exclude_match and exclude_match.group(0):
+ report.path_ignored(child, f"matches the --exclude regular expression")
+ continue
+
+ if child.is_dir():
+ yield from gen_python_files_in_dir(child, root, include, exclude, report)
+
+ elif child.is_file():
+ include_match = include.search(normalized_path)
+ if include_match:
+ yield child
+
+
+@lru_cache()
+def find_project_root(srcs: Iterable[str]) -> Path:
+ """Return a directory containing .git, .hg, or pyproject.toml.
+
+ That directory can be one of the directories passed in `srcs` or their
+ common parent.
+
+ If no directory in the tree contains a marker that would specify it's the
+ project root, the root of the file system is returned.
+ """
+ if not srcs:
+ return Path("/").resolve()
+
+ common_base = min(Path(src).resolve() for src in srcs)
+ if common_base.is_dir():
+ # Append a fake file so `parents` below returns `common_base_dir`, too.
+ common_base /= "fake-file"
+ for directory in common_base.parents:
+ if (directory / ".git").is_dir():
+ return directory
+
+ if (directory / ".hg").is_dir():
+ return directory
+
+ if (directory / "pyproject.toml").is_file():
+ return directory
+
+ return directory
+
+
+@dataclass
+class Report:
+ """Provides a reformatting counter. Can be rendered with `str(report)`."""
+
+ check: bool = False
+ quiet: bool = False
+ verbose: bool = False
+ change_count: int = 0
+ same_count: int = 0
+ failure_count: int = 0
+
+ def done(self, src: Path, changed: Changed) -> None:
+ """Increment the counter for successful reformatting. Write out a message."""
+ if changed is Changed.YES:
+ reformatted = "would reformat" if self.check else "reformatted"
+ if self.verbose or not self.quiet:
+ out(f"{reformatted} {src}")
+ self.change_count += 1
+ else:
+ if self.verbose:
+ if changed is Changed.NO:
+ msg = f"{src} already well formatted, good job."
+ else:
+ msg = f"{src} wasn't modified on disk since last run."
+ out(msg, bold=False)
+ self.same_count += 1
+
+ def failed(self, src: Path, message: str) -> None:
+ """Increment the counter for failed reformatting. Write out a message."""
+ err(f"error: cannot format {src}: {message}")
+ self.failure_count += 1
+
+ def path_ignored(self, path: Path, message: str) -> None:
+ if self.verbose:
+ out(f"{path} ignored: {message}", bold=False)
+
+ @property
+ def return_code(self) -> int:
+ """Return the exit code that the app should use.
+
+ This considers the current state of changed files and failures:
+ - if there were any failures, return 123;
+ - if any files were changed and --check is being used, return 1;
+ - otherwise return 0.
+ """
+ # According to http://tldp.org/LDP/abs/html/exitcodes.html starting with
+ # 126 we have special return codes reserved by the shell.
+ if self.failure_count:
+ return 123
+
+ elif self.change_count and self.check:
+ return 1