mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
ashton
186 lines
8.6 KiB
Python
186 lines
8.6 KiB
Python
from __future__ import annotations
|
|
|
|
import ast
|
|
|
|
|
|
class FStringLowerer(ast.NodeTransformer):
|
|
def visit_JoinedStr(self, node: ast.JoinedStr) -> ast.AST:
|
|
parts: list[ast.expr] = []
|
|
for value in node.values:
|
|
if isinstance(value, ast.Constant) and isinstance(value.value, str):
|
|
if value.value:
|
|
parts.append(ast.Constant(value=value.value))
|
|
elif isinstance(value, ast.FormattedValue):
|
|
parts.append(self._format_value(value))
|
|
else:
|
|
return self.generic_visit(node)
|
|
if not parts:
|
|
return ast.Constant(value="")
|
|
return ast.Call(
|
|
func=ast.Attribute(value=ast.Constant(value=""), attr="join", ctx=ast.Load()),
|
|
args=[ast.List(elts=parts, ctx=ast.Load())],
|
|
keywords=[],
|
|
)
|
|
|
|
def _format_value(self, node: ast.FormattedValue) -> ast.expr:
|
|
value = self.visit(node.value)
|
|
if node.conversion == 115:
|
|
value = ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[value], keywords=[])
|
|
elif node.conversion == 114:
|
|
value = ast.Call(func=ast.Name(id="repr", ctx=ast.Load()), args=[value], keywords=[])
|
|
elif node.conversion == 97:
|
|
value = ast.Call(func=ast.Name(id="ascii", ctx=ast.Load()), args=[value], keywords=[])
|
|
elif node.conversion != -1:
|
|
return ast.FormattedValue(value=value, conversion=node.conversion, format_spec=node.format_spec)
|
|
if node.format_spec is None:
|
|
spec: ast.expr = ast.Constant(value="")
|
|
else:
|
|
spec_node = self.visit(node.format_spec)
|
|
if not isinstance(spec_node, ast.expr):
|
|
return ast.FormattedValue(value=value, conversion=-1, format_spec=node.format_spec)
|
|
spec = spec_node
|
|
return ast.Call(func=ast.Name(id="format", ctx=ast.Load()), args=[value, spec], keywords=[])
|
|
|
|
|
|
class MatchLowerer(ast.NodeTransformer):
|
|
def __init__(self) -> None:
|
|
self._counter = 0
|
|
|
|
def visit_Match(self, node: ast.Match) -> ast.AST:
|
|
self.generic_visit(node)
|
|
if any(case.guard is not None for case in node.cases):
|
|
return node
|
|
subject_name = self._name()
|
|
subject_store = ast.Assign(targets=[ast.Name(id=subject_name, ctx=ast.Store())], value=node.subject)
|
|
chain = self._case_chain(subject_name, node.cases)
|
|
if chain is None:
|
|
return node
|
|
return [subject_store, chain]
|
|
|
|
def _name(self) -> str:
|
|
self._counter += 1
|
|
return f"_pw_match_{self._counter}"
|
|
|
|
def _case_chain(self, subject_name: str, cases: list[ast.match_case]) -> ast.stmt | None:
|
|
next_stmt: ast.stmt | None = None
|
|
for case in reversed(cases):
|
|
lowered = self._case_if(subject_name, case, [] if next_stmt is None else [next_stmt])
|
|
if lowered is None:
|
|
return None
|
|
next_stmt = lowered
|
|
return next_stmt
|
|
|
|
def _case_if(self, subject_name: str, case: ast.match_case, orelse: list[ast.stmt]) -> ast.stmt | None:
|
|
result = self._pattern(subject_name, case.pattern)
|
|
if result is None:
|
|
return None
|
|
test, binders = result
|
|
if case.guard is not None:
|
|
test = ast.BoolOp(op=ast.And(), values=[test, case.guard])
|
|
return ast.If(test=test, body=[*binders, *case.body], orelse=orelse)
|
|
|
|
def _pattern(self, subject_name: str, pattern: ast.pattern) -> tuple[ast.expr, list[ast.stmt]] | None:
|
|
subject = ast.Name(id=subject_name, ctx=ast.Load())
|
|
if isinstance(pattern, ast.MatchAs) and pattern.pattern is None:
|
|
if pattern.name is None:
|
|
return ast.Constant(value=True), []
|
|
return ast.Constant(value=True), [self._bind(pattern.name, subject)]
|
|
if isinstance(pattern, ast.MatchValue):
|
|
return ast.Compare(left=subject, ops=[ast.Eq()], comparators=[pattern.value]), []
|
|
if isinstance(pattern, ast.MatchSingleton):
|
|
return ast.Compare(left=subject, ops=[ast.Is()], comparators=[ast.Constant(value=pattern.value)]), []
|
|
if isinstance(pattern, ast.MatchSequence):
|
|
return self._sequence_pattern(subject_name, pattern)
|
|
return None
|
|
|
|
def _sequence_pattern(self, subject_name: str, pattern: ast.MatchSequence) -> tuple[ast.expr, list[ast.stmt]] | None:
|
|
patterns = pattern.patterns
|
|
star_index = next((index for index, item in enumerate(patterns) if isinstance(item, ast.MatchStar)), None)
|
|
tests: list[ast.expr] = [self._sequence_test(subject_name)]
|
|
binders: list[ast.stmt] = []
|
|
if star_index is None:
|
|
tests.append(self._len_compare(subject_name, ast.Eq(), len(patterns)))
|
|
else:
|
|
tests.append(self._len_compare(subject_name, ast.GtE(), len(patterns) - 1))
|
|
for index, item in enumerate(patterns):
|
|
if isinstance(item, ast.MatchAs):
|
|
value = self._index(subject_name, index, star_index, len(patterns))
|
|
if item.pattern is not None:
|
|
return None
|
|
if item.name is not None:
|
|
binders.append(self._bind(item.name, value))
|
|
elif isinstance(item, ast.MatchStar):
|
|
if item.name is not None:
|
|
binders.append(self._bind(item.name, self._star_slice(subject_name, index, len(patterns))))
|
|
else:
|
|
value = self._index(subject_name, index, star_index, len(patterns))
|
|
sub_name = self._name()
|
|
binders.append(ast.Assign(targets=[ast.Name(id=sub_name, ctx=ast.Store())], value=value))
|
|
sub_result = self._pattern(sub_name, item)
|
|
if sub_result is None:
|
|
return None
|
|
sub_test, sub_binders = sub_result
|
|
tests.append(sub_test)
|
|
binders.extend(sub_binders)
|
|
return ast.BoolOp(op=ast.And(), values=tests), binders
|
|
|
|
def _len_compare(self, subject_name: str, op: ast.cmpop, length: int) -> ast.expr:
|
|
return ast.Compare(
|
|
left=ast.Call(func=ast.Name(id="len", ctx=ast.Load()), args=[ast.Name(id=subject_name, ctx=ast.Load())], keywords=[]),
|
|
ops=[op],
|
|
comparators=[ast.Constant(value=length)],
|
|
)
|
|
|
|
def _sequence_test(self, subject_name: str) -> ast.expr:
|
|
subject = ast.Name(id=subject_name, ctx=ast.Load())
|
|
sequence_type = ast.Attribute(
|
|
value=ast.Attribute(
|
|
value=ast.Call(func=ast.Name(id="__import__", ctx=ast.Load()), args=[ast.Constant(value="collections.abc")], keywords=[]),
|
|
attr="abc",
|
|
ctx=ast.Load(),
|
|
),
|
|
attr="Sequence",
|
|
ctx=ast.Load(),
|
|
)
|
|
return ast.BoolOp(
|
|
op=ast.And(),
|
|
values=[
|
|
ast.Call(func=ast.Name(id="isinstance", ctx=ast.Load()), args=[subject, sequence_type], keywords=[]),
|
|
ast.UnaryOp(
|
|
op=ast.Not(),
|
|
operand=ast.Call(
|
|
func=ast.Name(id="isinstance", ctx=ast.Load()),
|
|
args=[
|
|
subject,
|
|
ast.Tuple(
|
|
elts=[
|
|
ast.Name(id="str", ctx=ast.Load()),
|
|
ast.Name(id="bytes", ctx=ast.Load()),
|
|
ast.Name(id="bytearray", ctx=ast.Load()),
|
|
],
|
|
ctx=ast.Load(),
|
|
),
|
|
],
|
|
keywords=[],
|
|
),
|
|
),
|
|
],
|
|
)
|
|
|
|
def _index(self, subject_name: str, index: int, star_index: int | None, total: int) -> ast.expr:
|
|
if star_index is not None and index > star_index:
|
|
index = index - total
|
|
return ast.Subscript(value=ast.Name(id=subject_name, ctx=ast.Load()), slice=ast.Constant(value=index), ctx=ast.Load())
|
|
|
|
def _star_slice(self, subject_name: str, index: int, total: int) -> ast.expr:
|
|
stop = index - total + 1
|
|
upper: ast.expr | None = None if stop == 0 else ast.Constant(value=stop)
|
|
return ast.Subscript(
|
|
value=ast.Call(func=ast.Name(id="list", ctx=ast.Load()), args=[ast.Name(id=subject_name, ctx=ast.Load())], keywords=[]),
|
|
slice=ast.Slice(lower=ast.Constant(value=index), upper=upper, step=None),
|
|
ctx=ast.Load(),
|
|
)
|
|
|
|
def _bind(self, name: str, value: ast.expr) -> ast.stmt:
|
|
return ast.Assign(targets=[ast.Name(id=name, ctx=ast.Store())], value=value)
|