import sys
from typing import (
- Collection,
Generic,
Iterator,
List,
Optional,
Set,
- Tuple,
TypeVar,
Union,
)
-if sys.version_info < (3, 8):
+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 import Final
+ 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
pygram.initialize(CACHE_DIR)
-syms = pygram.python_symbols
+syms: Final = pygram.python_symbols
# types
syms.with_stmt,
syms.funcdef,
syms.classdef,
+ syms.match_stmt,
+ syms.case_block,
}
STANDALONE_COMMENT: Final = 153
token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
syms.listmaker,
syms.testlist_gexp,
syms.testlist_star_expr,
+ syms.subject_expr,
+ syms.pattern,
}
TEST_DESCENDANTS: Final = {
syms.test,
"//=",
}
-IMPLICIT_TUPLE = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
-BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
-OPENING_BRACKETS = set(BRACKET.keys())
-CLOSING_BRACKETS = set(BRACKET.values())
-BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
-ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
+IMPLICIT_TUPLE: Final = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
+BRACKET: Final = {
+ token.LPAR: token.RPAR,
+ token.LSQB: token.RSQB,
+ token.LBRACE: token.RBRACE,
+}
+OPENING_BRACKETS: Final = set(BRACKET.keys())
+CLOSING_BRACKETS: Final = set(BRACKET.values())
+BRACKETS: Final = OPENING_BRACKETS | CLOSING_BRACKETS
+ALWAYS_NO_SPACE: Final = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
+
+RARROW = 55
+@mypyc_attr(allow_interpreted_subclasses=True)
class Visitor(Generic[T]):
"""Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
`complex_subscript` signals whether the given leaf is part of a subscription
which has non-trivial arguments, like arithmetic expressions or function calls.
"""
- NO = ""
- SPACE = " "
- DOUBLESPACE = " "
+ NO: Final = ""
+ SPACE: Final = " "
+ DOUBLESPACE: Final = " "
t = leaf.type
p = leaf.parent
v = leaf.value
):
return NO
- elif (
- prevp.type == token.RIGHTSHIFT
- 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
- ):
- # Python 2 print chevron
- return NO
elif prevp.type == token.AT and p.parent and p.parent.type == syms.decorator:
# no space in decorators
return NO
return NO
if not prev:
- if t == token.DOT:
- prevp = preceding_leaf(p)
- if not prevp or prevp.type != token.NUMBER:
- return NO
-
- elif t == token.LSQB:
+ if t == token.DOT or t == token.LSQB:
return NO
elif prev.type != token.COMMA:
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 = None
- last = 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:
return None
-def first_child_is_arith(node: Node) -> bool:
- """Whether first child is an arithmetic or a binary arithmetic expression"""
- expr_types = {
+def is_arith_like(node: LN) -> bool:
+ """Whether node is an arithmetic or a binary arithmetic expression"""
+ return node.type in {
syms.arith_expr,
syms.shift_expr,
syms.xor_expr,
syms.and_expr,
}
- return bool(node.children and node.children[0].type in expr_types)
def is_docstring(leaf: Leaf) -> 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:
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