X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/95e77cb5590a1499d3aa4cf7fe60481347191c35..935f303a0a7b794e722c7df00c906be285884874:/src/black/nodes.py
diff --git a/src/black/nodes.py b/src/black/nodes.py
index 7466670..edd201a 100644
--- a/src/black/nodes.py
+++ b/src/black/nodes.py
@@ -3,22 +3,8 @@ blib2to3 Node/Leaf transformation-related utility functions.
"""
import sys
-from typing import (
- Collection,
- Generic,
- Iterator,
- List,
- Optional,
- Set,
- Tuple,
- TypeVar,
- Union,
-)
-
-if sys.version_info >= (3, 8):
- from typing import Final
-else:
- from typing_extensions import Final
+from typing import Final, Generic, Iterator, List, Optional, Set, Tuple, TypeVar, Union
+
if sys.version_info >= (3, 10):
from typing import TypeGuard
else:
@@ -26,14 +12,12 @@ else:
from mypy_extensions import mypyc_attr
-# lib2to3 fork
-from blib2to3.pytree import Node, Leaf, type_repr, NL
-from blib2to3 import pygram
-from blib2to3.pgen2 import token
-
from black.cache import CACHE_DIR
+from black.mode import Mode, Preview
from black.strings import has_triple_quotes
-
+from blib2to3 import pygram
+from blib2to3.pgen2 import token
+from blib2to3.pytree import NL, Leaf, Node, type_repr
pygram.initialize(CACHE_DIR)
syms: Final = pygram.python_symbols
@@ -121,6 +105,7 @@ TEST_DESCENDANTS: Final = {
syms.term,
syms.power,
}
+TYPED_NAMES: Final = {syms.tname, syms.tname_star}
ASSIGNMENTS: Final = {
"=",
"+=",
@@ -187,15 +172,15 @@ class Visitor(Generic[T]):
yield from self.visit(child)
-def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
+def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str: # noqa: C901
"""Return whitespace prefix if needed for the given `leaf`.
`complex_subscript` signals whether the given leaf is part of a subscription
which has non-trivial arguments, like arithmetic expressions or function calls.
"""
- NO: Final = ""
- SPACE: Final = " "
- DOUBLESPACE: Final = " "
+ NO: Final[str] = ""
+ SPACE: Final[str] = " "
+ DOUBLESPACE: Final[str] = " "
t = leaf.type
p = leaf.parent
v = leaf.value
@@ -244,6 +229,14 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
# that, too.
return prevp.prefix
+ elif (
+ prevp.type == token.STAR
+ and parent_type(prevp) == syms.star_expr
+ and parent_type(prevp.parent) == syms.subscriptlist
+ ):
+ # No space between typevar tuples.
+ return NO
+
elif prevp.type in VARARGS_SPECIALS:
if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
return NO
@@ -282,7 +275,7 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
return NO
if t == token.EQUAL:
- if prev.type != syms.tname:
+ if prev.type not in TYPED_NAMES:
return NO
elif prev.type == token.EQUAL:
@@ -293,7 +286,7 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
elif prev.type != token.COMMA:
return NO
- elif p.type == syms.tname:
+ elif p.type in TYPED_NAMES:
# type names
if not prev:
prevp = preceding_leaf(p)
@@ -353,6 +346,11 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
return NO
+ elif Preview.walrus_subscript in mode and (
+ t == token.COLONEQUAL or prev.type == token.COLONEQUAL
+ ):
+ return SPACE
+
elif not complex_subscript:
return NO
@@ -402,6 +400,10 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901
elif p.type == syms.sliceop:
return NO
+ elif p.type == syms.except_clause:
+ if t == token.STAR:
+ return NO
+
return SPACE
@@ -439,27 +441,6 @@ def prev_siblings_are(node: Optional[LN], tokens: List[Optional[NodeType]]) -> b
return prev_siblings_are(node.prev_sibling, tokens[:-1])
-def last_two_except(leaves: List[Leaf], omit: Collection[LeafID]) -> Tuple[Leaf, Leaf]:
- """Return (penultimate, last) leaves skipping brackets in `omit` and contents."""
- stop_after: Optional[Leaf] = None
- last: Optional[Leaf] = None
- for leaf in reversed(leaves):
- if stop_after:
- if leaf is stop_after:
- stop_after = None
- continue
-
- if last:
- return leaf, last
-
- if id(leaf) in omit:
- stop_after = leaf.opening_bracket
- else:
- last = leaf
- else:
- raise LookupError("Last two leaves were also skipped")
-
-
def parent_type(node: Optional[LN]) -> Optional[NodeType]:
"""
Returns:
@@ -523,12 +504,14 @@ def container_of(leaf: Leaf) -> LN:
return container
-def first_leaf_column(node: Node) -> Optional[int]:
- """Returns the column of the first leaf child of a node."""
- for child in node.children:
- if isinstance(child, Leaf):
- return child.column
- return None
+def first_leaf_of(node: LN) -> Optional[Leaf]:
+ """Returns the first leaf of the node tree."""
+ if isinstance(node, Leaf):
+ return node
+ if node.children:
+ return first_leaf_of(node.children[0])
+ else:
+ return None
def is_arith_like(node: LN) -> bool:
@@ -582,9 +565,25 @@ def is_one_tuple(node: LN) -> bool:
)
-def is_one_tuple_between(opening: Leaf, closing: Leaf, leaves: List[Leaf]) -> bool:
- """Return True if content between `opening` and `closing` looks like a one-tuple."""
- if opening.type != token.LPAR and closing.type != token.RPAR:
+def is_tuple_containing_walrus(node: LN) -> bool:
+ """Return True if `node` holds a tuple that contains a walrus operator."""
+ if node.type != syms.atom:
+ return False
+ gexp = unwrap_singleton_parenthesis(node)
+ if gexp is None or gexp.type != syms.testlist_gexp:
+ return False
+
+ return any(child.type == syms.namedexpr_test for child in gexp.children)
+
+
+def is_one_sequence_between(
+ opening: Leaf,
+ closing: Leaf,
+ leaves: List[Leaf],
+ brackets: Tuple[int, int] = (token.LPAR, token.RPAR),
+) -> bool:
+ """Return True if content between `opening` and `closing` is a one-sequence."""
+ if (opening.type, closing.type) != brackets:
return False
depth = closing.bracket_depth + 1
@@ -721,6 +720,11 @@ def is_multiline_string(leaf: Leaf) -> bool:
def is_stub_suite(node: Node) -> bool:
"""Return True if `node` is a suite with a stub body."""
+
+ # If there is a comment, we want to keep it.
+ if node.prefix.strip():
+ return False
+
if (
len(node.children) != 4
or node.children[0].type != token.NEWLINE
@@ -729,6 +733,9 @@ def is_stub_suite(node: Node) -> bool:
):
return False
+ if node.children[3].prefix.strip():
+ return False
+
return is_stub_body(node.children[2])
@@ -742,7 +749,8 @@ def is_stub_body(node: LN) -> bool:
child = node.children[0]
return (
- child.type == syms.atom
+ not child.prefix.strip()
+ and child.type == syms.atom
and len(child.children) == 3
and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
)
@@ -792,12 +800,54 @@ def is_import(leaf: Leaf) -> bool:
)
-def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
- """Return True if the given leaf is a special comment.
- Only returns true for type comments for now."""
+def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
+ """Return True if the given leaf starts a with or async with statement."""
+ return bool(
+ leaf.type == token.NAME
+ and leaf.value == "with"
+ and leaf.parent
+ and leaf.parent.type == syms.with_stmt
+ ) or bool(
+ leaf.type == token.ASYNC
+ and leaf.next_sibling
+ and leaf.next_sibling.type == syms.with_stmt
+ )
+
+
+def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
+ """Return True if the given leaf starts an async def/for/with statement.
+
+ Note that `async def` can be either an `async_stmt` or `async_funcdef`,
+ the latter is used when it has decorators.
+ """
+ return bool(
+ leaf.type == token.ASYNC
+ and leaf.parent
+ and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
+ )
+
+
+def is_type_comment(leaf: Leaf) -> bool:
+ """Return True if the given leaf is a type comment. This function should only
+ be used for general type comments (excluding ignore annotations, which should
+ use `is_type_ignore_comment`). Note that general type comments are no longer
+ used in modern version of Python, this function may be deprecated in the future."""
+ t = leaf.type
+ v = leaf.value
+ return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:")
+
+
+def is_type_ignore_comment(leaf: Leaf) -> bool:
+ """Return True if the given leaf is a type comment with ignore annotation."""
t = leaf.type
v = leaf.value
- return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
+ return t in {token.COMMENT, STANDALONE_COMMENT} and is_type_ignore_comment_string(v)
+
+
+def is_type_ignore_comment_string(value: str) -> bool:
+ """Return True if the given string match with type comment with
+ ignore annotation."""
+ return value.startswith("# type: ignore")
def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
@@ -858,3 +908,19 @@ def is_rpar_token(nl: NL) -> TypeGuard[Leaf]:
def is_string_token(nl: NL) -> TypeGuard[Leaf]:
return nl.type == token.STRING
+
+
+def is_number_token(nl: NL) -> TypeGuard[Leaf]:
+ return nl.type == token.NUMBER
+
+
+def is_part_of_annotation(leaf: Leaf) -> bool:
+ """Returns whether this leaf is part of type annotations."""
+ ancestor = leaf.parent
+ while ancestor is not None:
+ if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
+ return True
+ if ancestor.parent and ancestor.parent.type == syms.tname:
+ return True
+ ancestor = ancestor.parent
+ return False