X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/136930fccb99320865622e12ffc21bdd45fd7501..668bace2aba1589aaa2bfd7c11787d79410bfd05:/src/black/nodes.py diff --git a/src/black/nodes.py b/src/black/nodes.py index 8bf1934..75a2347 100644 --- a/src/black/nodes.py +++ b/src/black/nodes.py @@ -19,11 +19,15 @@ if sys.version_info >= (3, 8): from typing import Final else: from typing_extensions import Final +if sys.version_info >= (3, 10): + from typing import TypeGuard +else: + from typing_extensions import TypeGuard from mypy_extensions import mypyc_attr # lib2to3 fork -from blib2to3.pytree import Node, Leaf, type_repr +from blib2to3.pytree import Node, Leaf, type_repr, NL from blib2to3 import pygram from blib2to3.pgen2 import token @@ -260,8 +264,8 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901 and prevp.parent and prevp.parent.type == syms.shift_expr and prevp.prev_sibling - and prevp.prev_sibling.type == token.NAME - and prevp.prev_sibling.value == "print" # type: ignore + and is_name_token(prevp.prev_sibling) + and prevp.prev_sibling.value == "print" ): # Python 2 print chevron return NO @@ -687,7 +691,7 @@ def is_yield(node: LN) -> bool: if node.type == syms.yield_expr: return True - if node.type == token.NAME and node.value == "yield": # type: ignore + if is_name_token(node) and node.value == "yield": return True if node.type != syms.atom: @@ -854,3 +858,19 @@ def ensure_visible(leaf: Leaf) -> None: leaf.value = "(" elif leaf.type == token.RPAR: leaf.value = ")" + + +def is_name_token(nl: NL) -> TypeGuard[Leaf]: + return nl.type == token.NAME + + +def is_lpar_token(nl: NL) -> TypeGuard[Leaf]: + return nl.type == token.LPAR + + +def is_rpar_token(nl: NL) -> TypeGuard[Leaf]: + return nl.type == token.RPAR + + +def is_string_token(nl: NL) -> TypeGuard[Leaf]: + return nl.type == token.STRING