Fix stmt lambda scoping

Resolves   #814.

DOCS.md

@@ -1728,7 +1728,7 @@ If the last `statement` (not followed by a semicolon) in a statement lambda is a
 
 Statement lambdas also support implicit lambda syntax such that `def => _` is equivalent to `def (_=None) => _` as well as explicitly marking them as pattern-matching such that `match def (x) => x` will be a pattern-matching function.
 
-Importantly, statement lambdas do not capture variables introduced only in the surrounding expression, e.g. inside of a list comprehension or normal lambda. To avoid such situations, only nest statement lambdas inside other statement lambdas, and explicitly partially apply a statement lambda to pass in a value from a list comprehension.
+Additionally, statement lambdas have slightly different scoping rules than normal lambdas. When a statement lambda is inside of an expression with an expression-local variable, such as a normal lambda or comprehension, the statement lambda will capture the value of the variable at the time that the statement lambda is defined (rather than a reference to the overall namespace as with normal lambdas). As a result, while `[=> y for y in range(2)] |> map$(call) |> list` is `[1, 1]`, `[def => y for y in range(2)] |> map$(call) |> list` is `[0, 1]`. Note that this only works for expression-local variables: to copy the entire namespace at the time of function definition, use [`copyclosure`](#copyclosure-functions) (which can be used with statement lambdas).
 
 Note that statement lambdas have a lower precedence than normal lambdas and thus capture things like trailing commas. To avoid confusion, statement lambdas should always be wrapped in their own set of parentheses.
 

coconut/compiler/grammar.py

@@ -802,6 +802,7 @@ class Grammar(object):
 
         refname = Forward()
         setname = Forward()
+        expr_setname = Forward()
         classname = Forward()
         name_ref = combine(Optional(backslash) + base_name)
         unsafe_name = combine(Optional(backslash.suppress()) + base_name)
@@ -955,13 +956,13 @@ class Grammar(object):
         expr = Forward()
         star_expr = Forward()
         dubstar_expr = Forward()
-        comp_for = Forward()
         test_no_cond = Forward()
         infix_op = Forward()
         namedexpr_test = Forward()
         # for namedexpr locations only supported in Python 3.10
         new_namedexpr_test = Forward()
-        lambdef = Forward()
+        comp_for = Forward()
+        comprehension_expr = Forward()
 
         typedef = Forward()
         typedef_default = Forward()
@@ -971,6 +972,10 @@ class Grammar(object):
         typedef_ellipsis = Forward()
         typedef_op_item = Forward()
 
+        expr_lambdef = Forward()
+        stmt_lambdef = Forward()
+        lambdef = expr_lambdef | stmt_lambdef
+
         negable_atom_item = condense(Optional(neg_minus) + atom_item)
 
         testlist = itemlist(test, comma, suppress_trailing=False)
@@ -1148,8 +1153,8 @@ class Grammar(object):
                 ZeroOrMore(
                     condense(
                         # everything here must end with setarg_comma
-                        setname + Optional(default) + setarg_comma
-                        | (star | dubstar) + setname + setarg_comma
+                        expr_setname + Optional(default) + setarg_comma
+                        | (star | dubstar) + expr_setname + setarg_comma
                         | star_sep_setarg
                         | slash_sep_setarg
                     )
@@ -1180,7 +1185,7 @@ class Grammar(object):
             # everything here must end with rparen
             rparen.suppress()
             | tokenlist(Group(call_item), comma) + rparen.suppress()
-            | Group(attach(addspace(test + comp_for), add_parens_handle)) + rparen.suppress()
+            | Group(attach(comprehension_expr, add_parens_handle)) + rparen.suppress()
             | Group(op_item) + rparen.suppress()
         )
         function_call = Forward()
@@ -1230,10 +1235,6 @@ class Grammar(object):
             comma,
         )
 
-        comprehension_expr = (
-            addspace(namedexpr_test + comp_for)
-            | invalid_syntax(star_expr + comp_for, "iterable unpacking cannot be used in comprehension")
-        )
         paren_atom = condense(lparen + any_of(
             # everything here must end with rparen
             rparen,
@@ -1282,7 +1283,7 @@ class Grammar(object):
         setmaker = Group(
             (new_namedexpr_test + FollowedBy(rbrace))("test")
             | (new_namedexpr_testlist_has_comma + FollowedBy(rbrace))("list")
-            | addspace(new_namedexpr_test + comp_for + FollowedBy(rbrace))("comp")
+            | (comprehension_expr + FollowedBy(rbrace))("comp")
             | (testlist_star_namedexpr + FollowedBy(rbrace))("testlist_star_expr")
         )
         set_literal_ref = lbrace.suppress() + setmaker + rbrace.suppress()
@@ -1382,6 +1383,9 @@ class Grammar(object):
         no_partial_trailer_atom_ref = atom + ZeroOrMore(no_partial_trailer)
         partial_atom_tokens = no_partial_trailer_atom + partial_trailer_tokens
 
+        # must be kept in sync with expr_assignlist block below
+        assignlist = Forward()
+        star_assign_item = Forward()
         simple_assign = Forward()
         simple_assign_ref = maybeparens(
             lparen,
@@ -1391,12 +1395,8 @@ class Grammar(object):
                 | setname
                 | passthrough_atom
             ),
-            rparen
+            rparen,
         )
-        simple_assignlist = maybeparens(lparen, itemlist(simple_assign, comma, suppress_trailing=False), rparen)
-
-        assignlist = Forward()
-        star_assign_item = Forward()
         base_assign_item = condense(
             simple_assign
             | lparen + assignlist + rparen
@@ -1406,6 +1406,30 @@ class Grammar(object):
         assign_item = base_assign_item | star_assign_item
         assignlist <<= itemlist(assign_item, comma, suppress_trailing=False)
 
+        # must be kept in sync with assignlist block above (but with expr_setname)
+        expr_assignlist = Forward()
+        expr_star_assign_item = Forward()
+        expr_simple_assign = Forward()
+        expr_simple_assign_ref = maybeparens(
+            lparen,
+            (
+                # refname if there's a trailer, expr_setname if not
+                (refname | passthrough_atom) + OneOrMore(ZeroOrMore(complex_trailer) + OneOrMore(simple_trailer))
+                | expr_setname
+                | passthrough_atom
+            ),
+            rparen,
+        )
+        expr_base_assign_item = condense(
+            expr_simple_assign
+            | lparen + expr_assignlist + rparen
+            | lbrack + expr_assignlist + rbrack
+        )
+        expr_star_assign_item_ref = condense(star + expr_base_assign_item)
+        expr_assign_item = expr_base_assign_item | expr_star_assign_item
+        expr_assignlist <<= itemlist(expr_assign_item, comma, suppress_trailing=False)
+
+        simple_assignlist = maybeparens(lparen, itemlist(simple_assign, comma, suppress_trailing=False), rparen)
         typed_assign_stmt = Forward()
         typed_assign_stmt_ref = simple_assign + colon.suppress() + typedef_test + Optional(equals.suppress() + test_expr)
         basic_stmt = addspace(ZeroOrMore(assignlist + equals) + test_expr)
@@ -1639,7 +1663,10 @@ class Grammar(object):
         unsafe_lambda_arrow = any_of(fat_arrow, arrow)
 
         keyword_lambdef_params = maybeparens(lparen, set_args_list, rparen)
-        arrow_lambdef_params = lparen.suppress() + set_args_list + rparen.suppress() | setname
+        arrow_lambdef_params = (
+            lparen.suppress() + set_args_list + rparen.suppress()
+            | expr_setname
+        )
 
         keyword_lambdef = Forward()
         keyword_lambdef_ref = addspace(keyword("lambda") + condense(keyword_lambdef_params + colon))
@@ -1651,7 +1678,6 @@ class Grammar(object):
             keyword_lambdef,
         )
 
-        stmt_lambdef = Forward()
         match_guard = Optional(keyword("if").suppress() + namedexpr_test)
         closing_stmt = longest(new_testlist_star_expr("tests"), unsafe_simple_stmt_item)
         stmt_lambdef_match_params = Group(lparen.suppress() + match_args_list + match_guard + rparen.suppress())
@@ -1698,8 +1724,9 @@ class Grammar(object):
             | fixto(always_match, "")
         )
 
-        lambdef <<= addspace(lambdef_base + test) | stmt_lambdef
-        lambdef_no_cond = addspace(lambdef_base + test_no_cond)
+        expr_lambdef_ref = addspace(lambdef_base + test)
+        lambdef_no_cond = Forward()
+        lambdef_no_cond_ref = addspace(lambdef_base + test_no_cond)
 
         typedef_callable_arg = Group(
             test("arg")
@@ -1808,11 +1835,15 @@ class Grammar(object):
             invalid_syntax(maybeparens(lparen, namedexpr, rparen), "PEP 572 disallows assignment expressions in comprehension iterable expressions")
             | test_item
         )
-        base_comp_for = addspace(keyword("for") + assignlist + keyword("in") + comp_it_item + Optional(comp_iter))
+        base_comp_for = addspace(keyword("for") + expr_assignlist + keyword("in") + comp_it_item + Optional(comp_iter))
         async_comp_for_ref = addspace(keyword("async") + base_comp_for)
         comp_for <<= base_comp_for | async_comp_for
         comp_if = addspace(keyword("if") + test_no_cond + Optional(comp_iter))
         comp_iter <<= any_of(comp_for, comp_if)
+        comprehension_expr_ref = (
+            addspace(namedexpr_test + comp_for)
+            | invalid_syntax(star_expr + comp_for, "iterable unpacking cannot be used in comprehension")
+        )
 
         return_stmt = addspace(keyword("return") - Optional(new_testlist_star_expr))
 
@@ -2547,7 +2578,7 @@ class Grammar(object):
         original_function_call_tokens = (
             lparen.suppress() + rparen.suppress()
             # we need to keep the parens here, since f(x for x in y) is fine but tail_call(f, x for x in y) is not
-            | condense(lparen + originalTextFor(test + comp_for) + rparen)
+            | condense(lparen + originalTextFor(comprehension_expr) + rparen)
             | attach(parens, strip_parens_handle)
         )
 

coconut/compiler/util.py

@@ -1166,7 +1166,7 @@ class Wrap(ParseElementEnhance):
     global_instance_counter = 0
     inside = False
 
-    def __init__(self, item, wrapper, greedy=False, include_in_packrat_context=False):
+    def __init__(self, item, wrapper, greedy=False, include_in_packrat_context=True):
         super(Wrap, self).__init__(item)
         self.wrapper = wrapper
         self.greedy = greedy
@@ -1225,10 +1225,14 @@ def __repr__(self):
         return self.wrapped_name
 
 
-def handle_and_manage(item, handler, manager):
+def manage(item, manager, greedy=True, include_in_packrat_context=False):
+    """Attach a manager to the given parse item."""
+    return Wrap(item, manager, greedy=greedy, include_in_packrat_context=include_in_packrat_context)
+
+
+def handle_and_manage(item, handler, manager, **kwargs):
     """Attach a handler and a manager to the given parse item."""
-    new_item = attach(item, handler)
-    return Wrap(new_item, manager, greedy=True)
+    return manage(attach(item, handler), manager, **kwargs)
 
 
 def disable_inside(item, *elems, **kwargs):

coconut/root.py

@@ -26,7 +26,7 @@
 VERSION = "3.0.4"
 VERSION_NAME = None
 # False for release, int >= 1 for develop
-DEVELOP = 20
+DEVELOP = 21
 ALPHA = False  # for pre releases rather than post releases
 
 assert DEVELOP is False or DEVELOP >= 1, "DEVELOP must be False or an int >= 1"

coconut/tests/src/cocotest/agnostic/primary_2.coco

@@ -444,6 +444,14 @@ def primary_test_2() -> bool:
     assert b"Abc" |> fmap$(.|32) == b"abc"
     assert bytearray(b"Abc") |> fmap$(.|32) == bytearray(b"abc")
     assert (bytearray(b"Abc") |> fmap$(.|32)) `isinstance` bytearray
+    assert 10 |> lift(+)((x -> x), (def y -> y)) == 20
+    assert (x -> def y -> (x, y))(1)(2) == (1, 2) == (x -> copyclosure def y -> (x, y))(1)(2)  # type: ignore
+    assert ((x, y) -> def z -> (x, y, z))(1, 2)(3) == (1, 2, 3) == (x -> y -> def z -> (x, y, z))(1)(2)(3)  # type: ignore
+    assert [def x -> (x, y) for y in range(10)] |> map$(call$(?, 10)) |> list == [(10, y) for y in range(10)]
+    assert [x -> (x, y) for y in range(10)] |> map$(call$(?, 10)) |> list == [(10, 9) for y in range(10)]
+    assert [=> y for y in range(2)] |> map$(call) |> list == [1, 1]
+    assert [def => y for y in range(2)] |> map$(call) |> list == [0, 1]
+    assert (x -> x -> def y -> (x, y))(1)(2)(3) == (2, 3)
 
     with process_map.multiple_sequential_calls():  # type: ignore
         assert map((+), range(3), range(4)$[:-1], strict=True) |> list == [0, 2, 4] == process_map((+), range(3), range(4)$[:-1], strict=True) |> list  # type: ignore

coconut/tests/src/cocotest/agnostic/util.coco

@@ -34,6 +34,14 @@ def assert_raises(c, exc):
     else:
         raise AssertionError(f"{c} failed to raise exception {exc}")
 
+def x `typed_eq` y = (type(x), x) == (type(y), y)
+
+def pickle_round_trip(obj) = (
+    obj
+    |> pickle.dumps
+    |> pickle.loads
+)
+
 try:
     prepattern()  # type: ignore
 except NameError, TypeError:
@@ -44,14 +52,6 @@ except NameError, TypeError:
             return addpattern(func, base_func, **kwargs)
         return pattern_prepender
 
-def x `typed_eq` y = (type(x), x) == (type(y), y)
-
-def pickle_round_trip(obj) = (
-    obj
-    |> pickle.dumps
-    |> pickle.loads
-)
-
 
 # Old functions:
 old_fmap = fmap$(starmap_over_mappings=True)

coconut/tests/src/cocotest/target_38/py38_test.coco

@@ -10,4 +10,6 @@ def py38_test() -> bool:
     assert 10 |> (x := .) == 10 == x
     assert 10 |> (x := .) |> (. + 1) == 11
     assert x == 10
+    assert not consume(y := i for i in range(10))
+    assert y == 9
     return True

coconut/tests/src/extras.coco

@@ -313,8 +313,14 @@ def g(x) = x
 
     assert parse("def f(x):\n ${var}", "xonsh") == "def f(x):\n    ${var}\n"
     assert "data ABC" not in parse("data ABC:\n ${var}", "xonsh")
-
     assert parse('"abc" "xyz"', "lenient") == "'abcxyz'"
+    assert "builder" not in parse("def x -> x", "lenient")
+    assert parse("def x -> x", "lenient").count("def") == 1
+    assert "builder" in parse("x -> def y -> (x, y)", "lenient")
+    assert parse("x -> def y -> (x, y)", "lenient").count("def") == 2
+    assert "builder" in parse("[def x -> (x, y) for y in range(10)]", "lenient")
+    assert parse("[def x -> (x, y) for y in range(10)]", "lenient").count("def") == 2
+    assert parse("123  # derp", "lenient") == "123  # derp"
 
     return True
 
@@ -465,6 +471,11 @@ async def async_map_test() =
 # Compiled Coconut: -----------------------------------------------------------
 
 type Num = int | float""".strip())
+    assert parse("type L[T] = list[T]").strip().endswith("""
+# Compiled Coconut: -----------------------------------------------------------
+
+_coconut_typevar_T_0 = _coconut.typing.TypeVar("_coconut_typevar_T_0")
+type L = list[_coconut_typevar_T_0]""".strip())
 
     setup(line_numbers=False, minify=True)
     assert parse("123  # derp", "lenient") == "123# derp"