# -*- coding: utf-8 -*-

"""Lambdas with multiple expressions, local variables, and a name."""

__all__ = ["multilambda",

"namedlambda",

"fn", "_",

"quicklambda",

"envify"]

from ast import (Lambda, Name, Assign, Subscript, Call, FunctionDef,

AsyncFunctionDef, Attribute, keyword, Dict, Constant, arg)

from copy import deepcopy

from mcpyrate.quotes import macros, q, u, n, a, h # noqa: F401

from mcpyrate import gensym

from mcpyrate.astcompat import NamedExpr

from mcpyrate.expander import MacroExpander

from mcpyrate.quotes import is_captured_value

from mcpyrate.splicing import splice_expression

from mcpyrate.utils import extract_bindings

from mcpyrate.walkers import ASTTransformer

from ..dynassign import dyn

from ..env import env

from ..misc import namelambda

from ..symbol import sym

from .letdo import _implicit_do, _do

from .letdoutil import islet, isenvassign, UnexpandedLetView, UnexpandedEnvAssignView, ExpandedDoView

from .nameutil import getname

from .util import (is_decorated_lambda, isx, has_deco,

destructure_decorated_lambda, detect_lambda)

# --------------------------------------------------------------------------------

# Macro interface

def multilambda(tree, *, syntax, expander, **kw):

"""[syntax, block] Supercharge your lambdas: multiple expressions, local variables.

For all ``lambda`` lexically inside the ``with multilambda`` block,

``[...]`` denotes a multiple-expression body with an implicit ``do``::

lambda ...: [expr0, ...] --> lambda ...: do[expr0, ...]

Only the outermost set of brackets around the body of a ``lambda`` denotes

a multi-expression body; the rest are interpreted as lists, as usual.

Examples::

with multilambda:

echo = lambda x: [print(x), x]

assert echo("hi there") == "hi there"

count = let[x << 0][

lambda: [x << x + 1,

x]]

assert count() == 1

assert count() == 2

mk12 = lambda: [[1, 2]]

assert mk12() == [1, 2]

For local variables, see ``do``.

"""

if syntax != "block":

raise SyntaxError("multilambda is a block macro only") # pragma: no cover

if syntax == "block" and kw['optional_vars'] is not None:

raise SyntaxError("multilambda does not take an as-part") # pragma: no cover

# Expand outside in.

# multilambda should expand first before any let[], do[] et al. that happen

# to be inside the block, to avoid misinterpreting implicit lambdas

# generated by those constructs.

with dyn.let(_macro_expander=expander): # implicit do (extra bracket notation) needs this.

return _multilambda(block_body=tree)

def namedlambda(tree, *, syntax, expander, **kw):

"""[syntax, block] Name lambdas implicitly.

Lexically inside a ``with namedlambda`` block, any literal ``lambda``

that is assigned to a name using one of the supported assignment forms

is named to have the name of the LHS of the assignment. The name is

captured at macro expansion time.

Naming modifies the original function object.

We support:

- Single-item assignments to a local name, ``f = lambda ...: ...``

- Named expressions (a.k.a. walrus operator, Python 3.8+),

``f := lambda ...: ...``

- Assignments to unpythonic environments, ``f << (lambda ...: ...)``

- Let bindings, ``let[[f << (lambda ...: ...)] in ...]``, using any

let syntax supported by unpythonic (here using the haskelly let-in

just as an example).

Support for other forms of assignment might or might not be added in a

future version.

Example::

with namedlambda:

f = lambda x: x**3 # assignment: name as "f"

let[x << 42, g << None, h << None][[

g << (lambda x: x**2), # env-assignment: name as "g"

h << f, # still "f" (no literal lambda on RHS)

(g(x), h(x))]]

foo = let[[f7 << (lambda x: x)] in f7] # let-binding: name as "f7"

The naming is performed using the function ``unpythonic.misc.namelambda``,

which will update ``__name__``, ``__qualname__`` and ``__code__.co_name``.

"""

if syntax != "block":

raise SyntaxError("namedlambda is a block macro only") # pragma: no cover

if syntax == "block" and kw['optional_vars'] is not None:

raise SyntaxError("namedlambda does not take an as-part") # pragma: no cover

# Two-pass macro. We pass in the expander to allow the macro to decide when to recurse.

with dyn.let(_macro_expander=expander):

return _namedlambda(block_body=tree)

def fn(tree, *, syntax, expander, **kw):

"""[syntax, expr] Underscore notation (quick lambdas) for Python.

Usage::

fn[body]

The ``fn[]`` macro creates a lambda. Each underscore in ``body``

introduces a new parameter.

Example::

func = fn[_ * _]

expands to::

func = lambda a0, a1: a0 * a1

The underscore is interpreted magically by ``fn[]``; but ``_`` itself

is not a macro, and has no special meaning outside ``fn[]``. The underscore

does **not** need to be imported for ``fn[]`` to recognize it.

But if you want to make your IDE happy, there is a symbol named ``_`` in

`unpythonic.syntax` you can import to silence any "undefined name" errors

regarding the use of ``_``. It is a regular run-time object, not a macro.

The macro does not descend into any nested ``fn[]``.

"""

if syntax != "expr":

raise SyntaxError("f is an expr macro only") # pragma: no cover

# What's my name in the current expander? (There may be several names.)

# https://github.com/Technologicat/mcpyrate/blob/master/doc/quasiquotes.md#hygienic-macro-recursion

bindings = extract_bindings(expander.bindings, fn)

mynames = list(bindings.keys())

return _fn(tree, mynames)

_ = sym("_") # for those who want to make their IDEs happy

def quicklambda(tree, *, syntax, expander, **kw):

"""[syntax, block] Make ``fn`` quick lambdas expand first.

To be able to transform correctly, the block macros in ``unpythonic.syntax``

that transform lambdas (e.g. ``multilambda``, ``tco``) need to see all

``lambda`` definitions written with Python's standard ``lambda``.

However, the ``fn`` macro uses the syntax ``f[...]``, which (to the analyzer)

does not look like a lambda definition. This macro changes the expansion

order, forcing any ``fn[...]`` lexically inside the block to expand before

any other macros do.

Any expression of the form ``fn[...]``, where ``fn`` is any name bound in the

current macro expander to the macro `unpythonic.syntax.fn`, is understood as

a quick lambda. (In plain English, this respects as-imports of the macro ``fn``.)

Example - a quick multilambda::

from unpythonic.syntax import macros, multilambda, quicklambda, fn, local

from unpythonic.syntax import _ # optional, makes IDEs happy

with quicklambda, multilambda:

func = fn[[local[x << _],

local[y << _],

x + y]]

assert func(1, 2) == 3

(This is of course rather silly, as an unnamed argument can only be mentioned

once. If we're giving names to them, a regular ``lambda`` is shorter to write.

The point is, this combo is now possible.)

"""

if syntax != "block":

raise SyntaxError("quicklambda is a block macro only") # pragma: no cover

if syntax == "block" and kw['optional_vars'] is not None:

raise SyntaxError("quicklambda does not take an as-part") # pragma: no cover

# This macro expands outside in.

#

# In `mcpyrate`, expander instances are cheap - so we create a second expander

# to which we register only the `f` macro, under whatever names it appears in

# the original expander. Thus it leaves all other macros alone. This is the

# official `mcpyrate` way to immediately expand only some particular macros

# inside the current macro invocation.

bindings = extract_bindings(expander.bindings, fn)

return MacroExpander(bindings, expander.filename).visit(tree)

def envify(tree, *, syntax, expander, **kw):

"""[syntax, block] Make formal parameters live in an unpythonic env.

The purpose is to allow overwriting formals using unpythonic's

expression-assignment ``name << value``. The price is that the references

to the arguments are copied into an env whenever an envified function is

entered.

Example - PG's accumulator puzzle (http://paulgraham.com/icad.html)::

with envify:

def foo(n):

return lambda i: n << n + i

Or even shorter::

with autoreturn, envify:

def foo(n):

lambda i: n << n + i

"""

if syntax != "block":

raise SyntaxError("envify is a block macro only") # pragma: no cover

if syntax == "block" and kw['optional_vars'] is not None:

raise SyntaxError("envify does not take an as-part") # pragma: no cover

# Two-pass macro.

with dyn.let(_macro_expander=expander):

return _envify(block_body=tree)

# --------------------------------------------------------------------------------

# Syntax transformers

def _multilambda(block_body):

class MultilambdaTransformer(ASTTransformer):

def transform(self, tree):

if is_captured_value(tree):

return tree # don't recurse!

if type(tree) is Lambda:

tree.body = _implicit_do(tree.body)

return self.generic_visit(tree)

# multilambda should expand first before any let[], do[] et al. that happen

# to be inside the block, to avoid misinterpreting implicit lambdas

# generated by those constructs.

return MultilambdaTransformer().visit(block_body)

def _namedlambda(block_body):

def issingleassign(tree):

return type(tree) is Assign and len(tree.targets) == 1 and type(tree.targets[0]) is Name

# detect a manual curry

def iscurrywithfinallambda(tree):

if not (type(tree) is Call and isx(tree.func, "curry") and tree.args):

return False

return type(tree.args[-1]) is Lambda

# Detect an autocurry from an already expanded "with autocurry".

# CAUTION: These must match what unpythonic.syntax.autocurry.autocurry uses in its output.

currycall_name = "currycall"

iscurryf = lambda name: name in ("curryf", "curry") # auto or manual curry in a "with autocurry"

def isautocurrywithfinallambda(tree):

# "currycall(..., curryf(lambda ...: ...))"

if not (type(tree) is Call and isx(tree.func, currycall_name) and tree.args and

type(tree.args[-1]) is Call and isx(tree.args[-1].func, iscurryf)):

return False

curryf_callnode = tree.args[-1]

lastarg = curryf_callnode.args[-1]

return type(lastarg) is Lambda

def iscallwithnamedargs(tree):

return type(tree) is Call and tree.keywords

# If `tree` is a (bare or decorated) lambda, inject run-time code to name

# it as `myname` (str); else return `tree` as-is.

def nameit(myname, tree):

match, thelambda = False, None

# For decorated lambdas, match any chain of one-argument calls.

# The `has_deco` check ignores any already named lambdas.

d = is_decorated_lambda(tree, mode="any") and not has_deco(["namelambda"], tree)

c = iscurrywithfinallambda(tree)

# This matches only during the second pass (after "with autocurry" has expanded)

# so it can't have namelambda already applied

if isautocurrywithfinallambda(tree): # "currycall(..., curryf(lambda ...: ...))"

match = True

thelambda = tree.args[-1].args[-1]

# --> "currycall(..., (namelambda(myname))(curryf(lambda ...: ...)))"

tree.args[-1].args[-1] = q[h[namelambda](u[myname])(a[thelambda])]

elif type(tree) is Lambda or d or c:

match = True

if d:

decorator_list, thelambda = destructure_decorated_lambda(tree)

elif c:

thelambda = tree.args[-1]

else:

thelambda = tree

tree = q[h[namelambda](u[myname])(a[tree])] # plonk it as outermost and hope for the best

return tree, thelambda, match

class NamedLambdaTransformer(ASTTransformer):

def transform(self, tree):

if is_captured_value(tree):

return tree # don't recurse!

if islet(tree, expanded=False): # let bindings

view = UnexpandedLetView(tree)

newbindings = []

for b in view.bindings:

b.elts[1], thelambda, match = nameit(getname(b.elts[0]), b.elts[1])

if match:

thelambda.body = self.visit(thelambda.body)

else:

b.elts[1] = self.visit(b.elts[1])

newbindings.append(b)

view.bindings = newbindings # write the new bindings (important!)

view.body = self.visit(view.body)

return tree

# assumption: no one left-shifts by a literal lambda :)

elif isenvassign(tree): # f << (lambda ...: ...)

view = UnexpandedEnvAssignView(tree)

view.value, thelambda, match = nameit(view.name, view.value)

if match:

thelambda.body = self.visit(thelambda.body)

else:

view.value = self.visit(view.value)

return tree

elif issingleassign(tree): # f = lambda ...: ...

tree.value, thelambda, match = nameit(getname(tree.targets[0]), tree.value)

if match:

thelambda.body = self.visit(thelambda.body)

else:

tree.value = self.visit(tree.value)

return tree

elif type(tree) is NamedExpr: # f := lambda ...: ... (Python 3.8+, added in unpythonic 0.15.0)

tree.value, thelambda, match = nameit(getname(tree.target), tree.value)

if match:

thelambda.body = self.visit(thelambda.body)

else:

tree.value = self.visit(tree.value)

return tree

elif iscallwithnamedargs(tree): # foo(f=lambda: ...)

for kw in tree.keywords:

if kw.arg is None: # **kwargs in Python 3.5+

kw.value = self.visit(kw.value)

continue

# a single named arg

kw.value, thelambda, match = nameit(kw.arg, kw.value)

if match:

thelambda.body = self.visit(thelambda.body)

else:

kw.value = self.visit(kw.value)

tree.args = self.visit(tree.args)

return tree

elif type(tree) is Dict: # {"f": lambda: ..., "g": lambda: ...}

lst = list(zip(tree.keys, tree.values))

for j in range(len(lst)):

k, v = tree.keys[j], tree.values[j]

if k is None: # {..., **d, ...}

tree.values[j] = self.visit(v)

else:

if type(k) is Constant:

thename = k.value

tree.values[j], thelambda, match = nameit(thename, v)

if match:

thelambda.body = self.visit(thelambda.body)

else:

tree.values[j] = self.visit(v)

else:

tree.keys[j] = self.visit(k)

tree.values[j] = self.visit(v)

return tree

return self.generic_visit(tree)

# outside in: transform in unexpanded let[] forms

newbody = NamedLambdaTransformer().visit(block_body)

newbody = dyn._macro_expander.visit_recursively(newbody)

# inside out: transform in expanded autocurry

newbody = NamedLambdaTransformer().visit(newbody)

# v0.15.0+: Finally, auto-name any still anonymous `lambda` with source location info.

# We must perform this in a separate pass so that expanded autocurry invocations

# are transformed correctly first.

class NamedLambdaFinalizationTransformer(ASTTransformer):

def transform(self, tree):

# Recurse into the lambda body in already named lambdas.

if is_decorated_lambda(tree, mode="any") and has_deco(["namelambda"], tree):

decorator_list, thelambda = destructure_decorated_lambda(tree)

thelambda.body = self.visit(thelambda.body)

return tree

elif type(tree) is Lambda:

if hasattr(tree, "lineno"):

thename = f"<lambda at {dyn._macro_expander.filename}:{tree.lineno}>"

tree, thelambda, match = nameit(thename, tree)

if match:

thelambda.body = self.visit(thelambda.body)

else:

tree = self.visit(tree)

return tree

return self.generic_visit(tree)

return NamedLambdaFinalizationTransformer().visit(newbody)

# The function `f` is adapted from the `f` macro in `macropy.quick_lambda`,

# stripped into a bare syntax transformer., and then the `@Walker` inside

# converted to a `mcpyrate` `ASTTransformer`. We have also added the code

# to ignore any nested `f[]`.

#

# Used under the MIT license.

# Copyright (c) 2013-2018, Li Haoyi, Justin Holmgren, Alberto Berti and all the other contributors.

def _fn(tree, mynames=()):

class UnderscoreTransformer(ASTTransformer):

def transform(self, tree):

if is_captured_value(tree):

return tree # don't recurse!

# Don't recurse into nested `f[]`.

# TODO: This would benefit from macro destructuring in the expander.

# TODO: See https://github.com/Technologicat/mcpyrate/issues/3

if type(tree) is Subscript and type(tree.value) is Name and tree.value.id in mynames:

return tree

elif type(tree) is Name and tree.id == "_":

name = gensym("_")

tree.id = name

self.collect(name)

return self.generic_visit(tree)

ut = UnderscoreTransformer()

tree = ut.visit(tree)

used_names = ut.collected

tree = q[lambda _: a[tree]] # noqa: F811, it's a placeholder overwritten at the next line.

tree.args.args = [arg(arg=x) for x in used_names]

return tree

def _envify(block_body):

# first pass, outside-in

userlambdas = detect_lambda(block_body)

# Expand inside-out to easily support lexical scoping.

block_body = dyn._macro_expander.visit_recursively(block_body)

# second pass, inside-out

def getargs(tree): # tree: FunctionDef, AsyncFunctionDef, Lambda

a = tree.args

allargs = a.posonlyargs + a.args + a.kwonlyargs

argnames = [x.arg for x in allargs]

if a.vararg:

argnames.append(a.vararg.arg)

if a.kwarg:

argnames.append(a.kwarg.arg)

return argnames

def isfunctionoruserlambda(tree):

return ((type(tree) in (FunctionDef, AsyncFunctionDef)) or

(type(tree) is Lambda and id(tree) in userlambdas))

# Create a renamed reference to the env() constructor to be sure the Call

# nodes added by us have a unique .func (not used by other macros or user code)

_envify = env

class EnvifyTransformer(ASTTransformer):

def transform(self, tree):

if is_captured_value(tree):

return tree # don't recurse!

bindings = self.state.bindings

enames = self.state.enames

def isourupdate(thecall):

if type(thecall.func) is not Attribute:

return False

return thecall.func.attr == "update" and any(isx(thecall.func.value, x) for x in enames)

if isfunctionoruserlambda(tree):

argnames = getargs(tree)

if argnames:

# prepend env init to function body, update bindings

kws = [keyword(arg=k, value=q[n[k]]) for k in argnames] # "x" --> x

newbindings = bindings.copy()

if type(tree) in (FunctionDef, AsyncFunctionDef):

ename = gensym("e")

theenv = q[h[_envify]()]

theenv.keywords = kws

with q as quoted:

n[ename] = a[theenv]

assignment = quoted[0]

tree.body.insert(0, assignment)

elif type(tree) is Lambda and id(tree) in userlambdas:

# We must in general inject a new do[] even if one is already there,

# due to scoping rules. If the user code writes to the same names in

# its do[] env, this shadows the formals; if it then pops one of its names,

# the name should revert to mean the formal parameter.

#

# inject a do[] and reuse its env

tree.body = _do(q[n["_here_"],

a[tree.body]])

view = ExpandedDoView(tree.body) # view.body: [(lambda e14: ...), ...]

ename = view.body[0].args.args[0].arg # do[] environment name

theupdate = q[n[f"{ename}.update"]]

thecall = q[a[theupdate]()]

thecall.keywords = kws

tree.body = splice_expression(thecall, tree.body, "_here_")

newbindings.update({k: q[n[f"{ename}.{k}"]] for k in argnames}) # "x" --> e.x

self.generic_withstate(tree, enames=(enames + [ename]), bindings=newbindings)

else:

# leave alone the _envify() added by us

if type(tree) is Call and (isx(tree.func, "_envify") or isourupdate(tree)):

# don't recurse

return tree

# transform env-assignments into our envs

elif isenvassign(tree):

view = UnexpandedEnvAssignView(tree)

if view.name in bindings.keys():

# Grab the envname from the actual binding of "varname", of the form `e.varname`

# (so it's the `id` of a `Name` that is the `value` of an `Attribute`).

envset = q[n[f"{bindings[view.name].value.id}.set"]]

newvalue = self.visit(view.value)

return q[a[envset](u[view.name], a[newvalue])]

# transform references to currently active bindings

# x --> e14.x

# It doesn't matter if this hits an already expanded inner `with envify`,

# because the gensymmed environment name won't be in our bindings, and the "x"

# has become the `attr` in an `Attribute` node.

elif type(tree) is Name and tree.id in bindings.keys():

# Preserve the Load/Store/Del context of the name.

ctx = getattr(tree, "ctx", None)

out = deepcopy(bindings[tree.id])

out.ctx = ctx

return out

return self.generic_visit(tree)

return EnvifyTransformer(bindings={}, enames=[]).visit(block_body)