265 lines
9.6 KiB
Python
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)
|