KiBot/kibot/mcpyrate/splicing.py

265 lines
9.6 KiB
Python

# -*- coding: utf-8; -*-
"""Utilities for splicing the actual code into a code template."""
__all__ = ["splice_expression", "splice_statements", "splice_dialect"]
import ast
from copy import deepcopy
from .astfixers import fix_locations
from .coreutils import ismacroimport
from .markers import ASTMarker
from .utils import getdocstring
from .walkers import ASTTransformer
def splice_expression(expr, template, tag="__paste_here__"):
"""Splice `expr` into `template`.
This is somewhat like `mcpyrate.quotes.a`, but must be called from outside the
quoted snippet.
Parameters:
`expr`: an expression AST node, or an AST marker containing such.
The expression you would like to splice in.
`template`: an AST node, or a `list` of AST nodes.
Template into which to splice `expr`.
Must contain a paste-here indicator AST node that specifies where
`expr` is to be spliced in. The node is expected to have the format::
ast.Name(id=tag)
Or in plain English, it's a bare identifier.
The first-found instance (in AST scan order) of the paste-here indicator
is replaced by `expr`.
If the paste-here indicator appears multiple times, second and further
instances are replaced with a `copy.deepcopy` of `expr` so that they will
stay independent during any further AST edits.
`tag`: `str`
The name of the paste-here indicator in `template`.
Returns `template` with `expr` spliced in. Note `template` is **not** copied,
and will be mutated in-place.
"""
if not template:
return expr
if not (isinstance(expr, ast.expr) or
(isinstance(expr, ASTMarker) and isinstance(expr.body, ast.expr))):
raise TypeError(f"`expr` must be an expression AST node or AST marker containing such; got {type(expr)} with value {repr(expr)}")
if not isinstance(template, (ast.AST, list)):
raise TypeError(f"`template` must be an AST or `list`; got {type(template)} with value {repr(template)}")
def ispastehere(tree):
return type(tree) is ast.Expr and type(tree.value) is ast.Name and tree.value.id == tag
class ExpressionSplicer(ASTTransformer):
def __init__(self):
self.first = True
super().__init__()
def transform(self, tree):
if ispastehere(tree):
if not self.first:
return deepcopy(expr)
self.first = False
return expr
return self.generic_visit(tree)
return ExpressionSplicer().visit(template)
# TODO: this is actually a generic list-of-ast splicer, not specific to statements.
def splice_statements(body, template, tag="__paste_here__"):
"""Splice `body` into `template`.
This is somewhat like `mcpyrate.quotes.a`, but must be called from outside the
quoted snippet.
Parameters:
`body`: `list` of statements
The statements you would like to splice in.
`template`: `list` of statements
Template into which to splice `body`.
Must contain a paste-here indicator AST node, in a statement position,
that specifies where `body` is to be spliced in. The node is expected
to have the format::
ast.Expr(value=ast.Name(id=tag))
Or in plain English, it's a bare identifier in a statement position.
The first-found instance (in AST scan order) of the paste-here indicator
is replaced by `body`.
If the paste-here indicator appears multiple times, second and further
instances are replaced with a `copy.deepcopy` of `body` so that they will
stay independent during any further AST edits.
`tag`: `str`
The name of the paste-here indicator in `template`.
Returns `template` with `body` spliced in. Note `template` is **not** copied,
and will be mutated in-place.
Example::
from mcpyrate.quotes import macros, q
from mcpyrate.splicing import splice_statements
body = [...] # a list of statements
with q as template:
...
__paste_here__
...
splice_statements(body, template)
(Flake8 will complain about the undefined name `__paste_here__`. You can silence
it with the appropriate `# noqa`, or to make it happy, import the `n` macro from
`mcpyrate.quotes` and use `n["__paste_here__"]` instead of a plain `__paste_here__`.)
"""
if isinstance(body, ast.AST):
body = [body]
if isinstance(template, ast.AST):
body = [template]
if not body:
raise ValueError("expected at least one statement in `body`")
if not template:
return body
def ispastehere(tree):
return type(tree) is ast.Expr and type(tree.value) is ast.Name and tree.value.id == tag
class StatementSplicer(ASTTransformer):
def __init__(self):
self.first = True
super().__init__()
def transform(self, tree):
if ispastehere(tree):
if not self.first:
return deepcopy(body)
self.first = False
return body
return self.generic_visit(tree)
return StatementSplicer().visit(template)
def splice_dialect(body, template, tag="__paste_here__"):
"""In a dialect AST transformer, splice module `body` into `template`.
On top of what `splice_statements` does, this function handles macro-imports
and dialect-imports specially, gathering them all at the top level of the
final module body, so that mcpyrate sees them when the module is sent to
the macro expander.
Any dialect-imports in the template are placed first (in the order they
appear in the template), followed by any dialect-imports in the user code
(in the order they appear in the user code), followed by macro-imports in
the template, then macro-imports in the user code.
This also handles the module docstring and the magic `__all__` (if any)
from `body`. The docstring comes first, before dialect-imports. The magic
`__all__` is placed after dialect-imports, before macro-imports.
Parameters:
`body`: `list` of statements
Original module body from the user code (input).
`template`: `list` of statements
Template for the final module body (output).
Must contain a paste-here indicator as in `splice_statements`.
`tag`: `str`
The name of the paste-here indicator in `template`.
Returns `template` with `body` spliced in. Note `template` is **not** copied,
and will be mutated in-place.
Also `body` is mutated, to remove macro-imports, `__all__` and the module
docstring; these are pasted into the final result.
"""
if isinstance(body, ast.AST):
body = [body]
if isinstance(template, ast.AST):
body = [template]
if not body:
raise ValueError("expected at least one statement in `body`")
if not template:
return body
# Generally speaking, dialect templates are fully macro-generated
# quasiquoted snippets with no source location info to start with.
# Even if they have location info, it's for a different file compared
# to the use site where `body` comes from.
#
# Pretend the template code appears at the beginning of the user module.
#
# TODO: It would be better to pretend it appears at the line that has the dialect-import.
# TODO: Requires a `lineno` parameter here, and `DialectExpander` must be modified to supply it.
# TODO: We could extract the `lineno` in `find_dialectimport_ast` and then pass it to the
# TODO: user-defined dialect AST transformer, so it could pass it to us if it wants to.
for stmt in template:
fix_locations(stmt, body[0], mode="overwrite")
if getdocstring(body):
docstring, *body = body
docstring = [docstring]
else:
docstring = []
def extract_magic_all(tree):
def ismagicall(tree):
if not (type(tree) is ast.Assign and len(tree.targets) == 1):
return False
target = tree.targets[0]
return type(target) is ast.Name and target.id == "__all__"
class MagicAllExtractor(ASTTransformer):
def transform(self, tree):
if ismagicall(tree):
self.collect(tree)
return None
# We get just the top level of body by not recursing.
return tree
w = MagicAllExtractor()
w.visit(tree)
return tree, w.collected
body, user_magic_all = extract_magic_all(body)
def extract_macroimports(tree, *, magicname="macros"):
class MacroImportExtractor(ASTTransformer):
def transform(self, tree):
if ismacroimport(tree, magicname):
self.collect(tree)
return None
return self.generic_visit(tree)
w = MacroImportExtractor()
w.visit(tree)
return tree, w.collected
template, template_dialect_imports = extract_macroimports(template, magicname="dialects")
template, template_macro_imports = extract_macroimports(template)
body, user_dialect_imports = extract_macroimports(body, magicname="dialects")
body, user_macro_imports = extract_macroimports(body)
finalbody = splice_statements(body, template, tag)
return (docstring +
user_magic_all +
template_dialect_imports + user_dialect_imports +
template_macro_imports + user_macro_imports +
finalbody)