From cc066dfa4b69744e8999faac023ed10a3db34f69 Mon Sep 17 00:00:00 2001
From: pdeffebach <23196228+pdeffebach@users.noreply.github.com>
Date: Fri, 24 Sep 2021 17:36:01 -0400
Subject: [PATCH] Another attempt at an astable flag (#298)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* initial attempt

* finally working

* start adding tests

* more tests

* more tests

* add docstring

* tests pass

* add ByRow in docstring

* add type annotation

* better docs

* more docs fixes

* update index.md

* Apply suggestions from code review

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>

* clean named tuple creation

* add example with string

* grouping tests

* Update src/macros.jl

Co-authored-by: Bogumił Kamiński <bkamins@sgh.waw.pl>

* changes

* fix some errors

* add macro check

* add errors for bad flag combo

* better grouping tests

* Update src/parsing_astable.jl

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>

* add snipper to transform, select, combine, by

* add mutating tests

* get rid of debugging printin

* Apply suggestions from code review

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>
Co-authored-by: Bogumił Kamiński <bkamins@sgh.waw.pl>
---
 Project.toml                 |   3 +-
 docs/src/index.md            |  32 +++++-
 src/DataFramesMeta.jl        |   5 +-
 src/macros.jl                | 188 ++++++++++++++++++++++++++----
 src/parsing.jl               |  17 ++-
 src/parsing_astable.jl       | 107 +++++++++++++++++
 test/astable_flag.jl         | 216 +++++++++++++++++++++++++++++++++++
 test/deprecated.jl           |   3 -
 test/function_compilation.jl |   4 +-
 test/grouping.jl             |   6 -
 test/runtests.jl             |   1 +
 11 files changed, 539 insertions(+), 43 deletions(-)
 create mode 100644 src/parsing_astable.jl
 create mode 100644 test/astable_flag.jl

diff --git a/Project.toml b/Project.toml
index fbb357d4..fbebc46e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,14 +6,15 @@ version = "0.9.1"
 Chain = "8be319e6-bccf-4806-a6f7-6fae938471bc"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [compat]
+Chain = "0.4"
 DataFrames = "1"
 MacroTools = "0.5"
 Reexport = "0.2, 1"
 julia = "1"
-Chain = "0.4"
 
 [extras]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
diff --git a/docs/src/index.md b/docs/src/index.md
index b110d01e..2ac92866 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -22,6 +22,7 @@ In addition, DataFramesMeta provides
   convenient syntax. 
 * `@byrow` for applying functions to each row of a data frame (only supported inside other macros).
 * `@passmissing` for propagating missing values inside row-wise DataFramesMeta.jl transformations.
+* `@astable` to create multiple columns within a single transformation.
 * `@chain`, from [Chain.jl](https://github.com/jkrumbiegel/Chain.jl) for piping the above macros together, similar to [magrittr](https://cran.r-project.org/web/packages/magrittr/vignettes/magrittr.html)'s
   `%>%` in R. 
 
@@ -396,11 +397,38 @@ julia> @rtransform df @passmissing x = parse(Int, :x_str)
    3 │ missing  missing
 ```
 
+## Creating multiple columns at once with `@astable`
+
+Often new variables may depend on the same intermediate calculations. `@astable` makes it easy to create multiple
+new variables in the same operation, yet have them share
+information. 
+
+In a single block, all assignments of the form `:y = f(:x)` 
+or `$y = f(:x)` at the top-level generate new columns. In the second example, `y`
+must be a string or `Symbol`. 
+
+```
+julia> df = DataFrame(a = [1, 2, 3], b = [400, 500, 600]);
+
+julia> @transform df @astable begin 
+           ex = extrema(:b)
+           :b_first = :b .- first(ex)
+           :b_last = :b .- last(ex)
+       end
+3×4 DataFrame
+ Row │ a      b      b_first  b_last 
+     │ Int64  Int64  Int64    Int64  
+─────┼───────────────────────────────
+   1 │     1    400        0    -200
+   2 │     2    500      100    -100
+   3 │     3    600      200       0
+```
+
+
 ## [Working with column names programmatically with `$`](@id dollar)
 
 DataFramesMeta provides the special syntax `$` for referring to 
-columns in a data frame via a `Symbol`, string, or column position as either
-a literal or a variable. 
+columns in a data frame via a `Symbol`, string, or column position as either a literal or a variable. 
 
 ```julia
 df = DataFrame(A = 1:3, B = [2, 1, 2])
diff --git a/src/DataFramesMeta.jl b/src/DataFramesMeta.jl
index fec20a35..56914e42 100644
--- a/src/DataFramesMeta.jl
+++ b/src/DataFramesMeta.jl
@@ -4,6 +4,8 @@ using Reexport
 
 using MacroTools
 
+using OrderedCollections: OrderedCollections
+
 @reexport using DataFrames
 
 @reexport using Chain
@@ -16,12 +18,13 @@ export @with,
        @transform, @select, @transform!, @select!,
        @rtransform, @rselect, @rtransform!, @rselect!,
        @eachrow, @eachrow!,
-       @byrow, @passmissing,
+       @byrow, @passmissing, @astable,
        @based_on, @where # deprecated
 
 const DOLLAR = raw"$"
 
 include("parsing.jl")
+include("parsing_astable.jl")
 include("macros.jl")
 include("linqmacro.jl")
 include("eachrow.jl")
diff --git a/src/macros.jl b/src/macros.jl
index e954a371..7fbfad8f 100644
--- a/src/macros.jl
+++ b/src/macros.jl
@@ -282,11 +282,10 @@ macro byrow(args...)
     throw(ArgumentError("@byrow is deprecated outside of DataFramesMeta macros."))
 end
 
-
 """
-    passmissing(args...)
+    @passmissing(args...)
 
-Propograte missing values inside DataFramesMeta.jl macros.
+Propagate missing values inside DataFramesMeta.jl macros.
 
 
 `@passmissing` is not a "real" Julia macro but rather serves as a "flag"
@@ -350,6 +349,156 @@ macro passmissing(args...)
     throw(ArgumentError("@passmissing only works inside DataFramesMeta macros."))
 end
 
+const astable_docstring_snippet = """
+    Transformations can also use the macro-flag [`@astable`](@ref) for creating multiple
+    new columns at once and letting transformations share the same name-space.
+    See `? @astable` for more details.
+    """
+
+"""
+    @astable(args...)
+
+Return a `NamedTuple` from a single transformation inside the DataFramesMeta.jl
+macros, `@select`, `@transform`, and their mutating and row-wise equivalents.
+
+`@astable` acts on a single block. It works through all top-level expressions
+and collects all such expressions of the form `:y = ...` or `$(DOLLAR)y = ...`, i.e. assignments to a
+`Symbol` or an escaped column identifier, which is a syntax error outside of
+DataFramesMeta.jl macros. At the end of the expression, all assignments are collected
+into a `NamedTuple` to be used with the `AsTable` destination in the DataFrames.jl
+transformation mini-language.
+
+Concretely, the expressions
+
+```
+df = DataFrame(a = 1)
+
+@rtransform df @astable begin
+    :x = 1
+    y = 50
+    :z = :x + y + :a
+end
+```
+
+become the pair
+
+```
+function f(a)
+    x_t = 1
+    y = 50
+    z_t = x_t + y + a
+
+    (; x = x_t, z = z_t)
+end
+
+transform(df, [:a] => ByRow(f) => AsTable)
+```
+
+`@astable` has two major advantages at the cost of increasing complexity.
+First, `@astable` makes it easy to create multiple columns from a single
+transformation, which share a scope. For example, `@astable` allows
+for the following (where `:x` and `:x_2` exist in the data frame already).
+
+```
+@transform df @astable begin
+    m = mean(:x)
+    :x_demeaned = :x .- m
+    :x2_demeaned = :x2 .- m
+end
+```
+
+The creation of `:x_demeaned` and `:x2_demeaned` both share the variable `m`,
+which does not need to be calculated twice.
+
+Second, `@astable` is useful when performing intermediate calculations
+and storing their results in new columns. For example, the following fails.
+
+```
+@rtransform df begin
+    :new_col_1 = :x + :y
+    :new_col_2 = :new_col_1 + :z
+end
+```
+
+This because DataFrames.jl does not guarantee sequential evaluation of
+transformations. `@astable` solves this problem
+
+@rtransform df @astable begin
+    :new_col_1 = :x + :y
+    :new_col_2 = :new_col_1 + :z
+end
+
+Column assignment in `@astable` follows similar rules as
+column assignment in other DataFramesMeta.jl macros. The left-
+-hand-side of a column assignment can be either a `Symbol` or any
+expression which evaluates to a `Symbol` or `AbstractString`. For example
+`:y = ...`, and `$(DOLLAR)y = ...` are both valid ways of assigning a new column.
+However unlike other DataFramesMeta.jl macros, multi-column assignments via
+`AsTable` are disallowed. The following will fail.
+
+```
+@transform df @astable begin
+    $AsTable = :x
+end
+```
+
+References to existing columns also follow the same
+rules as other DataFramesMeta.jl macros.
+
+### Examples
+
+```
+julia> df = DataFrame(a = [1, 2, 3], b = [4, 5, 6]);
+
+julia> d = @rtransform df @astable begin
+           :x = 1
+           y = 5
+           :z = :x + y
+       end
+3×4 DataFrame
+ Row │ a      b      x      z
+     │ Int64  Int64  Int64  Int64
+─────┼────────────────────────────
+   1 │     1      4      1      6
+   2 │     2      5      1      6
+   3 │     3      6      1      6
+
+julia> df = DataFrame(a = [1, 1, 2, 2], b = [5, 6, 70, 80]);
+
+julia> @by df :a @astable begin
+            ex = extrema(:b)
+            :min_b = first(ex)
+            :max_b = last(ex)
+       end
+2×3 DataFrame
+ Row │ a      min_b  max_b
+     │ Int64  Int64  Int64
+─────┼─────────────────────
+   1 │     1      5      6
+   2 │     2     70     80
+
+julia> new_col = "New Column";
+
+julia> @rtransform df @astable begin
+           f_a = first(:a)
+           $(DOLLAR)new_col = :a + :b + f_a
+           :y = :a * :b
+       end
+4×4 DataFrame
+ Row │ a      b      New Column  y
+     │ Int64  Int64  Int64       Int64
+─────┼─────────────────────────────────
+   1 │     1      5           7      5
+   2 │     1      6           8      6
+   3 │     2     70          74    140
+   4 │     2     80          84    160
+```
+
+"""
+macro astable(args...)
+    throw(ArgumentError("@astable only works inside DataFramesMeta macros."))
+end
+
 ##############################################################################
 ##
 ## @with
@@ -1097,6 +1246,8 @@ transformations by row, `@transform` allows `@byrow` at the
 beginning of a block of transformations (i.e. `@byrow begin... end`).
 All transformations in the block will operate by row.
 
+$astable_docstring_snippet
+
 ### Examples
 
 ```jldoctest
@@ -1233,6 +1384,8 @@ transform!ations by row, `@transform!` allows `@byrow` at the
 beginning of a block of transform!ations (i.e. `@byrow begin... end`).
 All transform!ations in the block will operate by row.
 
+$astable_docstring_snippet
+
 ### Examples
 
 ```jldoctest
@@ -1345,6 +1498,8 @@ transformations by row, `@select` allows `@byrow` at the
 beginning of a block of selectations (i.e. `@byrow begin... end`).
 All transformations in the block will operate by row.
 
+$astable_docstring_snippet
+
 ### Examples
 
 ```jldoctest
@@ -1465,6 +1620,8 @@ transformations by row, `@select!` allows `@byrow` at the
 beginning of a block of select!ations (i.e. `@byrow begin... end`).
 All transformations in the block will operate by row.
 
+$astable_docstring_snippet
+
 ### Examples
 
 ```jldoctest
@@ -1546,17 +1703,6 @@ function combine_helper(x, args...; deprecation_warning = false)
 
     exprs, outer_flags = create_args_vector(args...)
 
-    fe = first(exprs)
-    if length(exprs) == 1 &&
-        get_column_expr(fe) === nothing &&
-        !(fe.head == :(=) || fe.head == :kw)
-
-        @warn "Returning a Table object from @by and @combine now requires `$(DOLLAR)AsTable` on the LHS."
-
-        lhs = Expr(:$, :AsTable)
-        exprs = ((:($lhs = $fe)),)
-    end
-
     t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags) for ex in exprs)
 
     quote
@@ -1592,6 +1738,8 @@ and
 @combine(df, :mx = mean(:x), :sx = std(:x))
 ```
 
+$astable_docstring_snippet
+
 ### Examples
 
 ```julia
@@ -1666,16 +1814,6 @@ end
 function by_helper(x, what, args...)
     # Only allow one argument when returning a Table object
     exprs, outer_flags = create_args_vector(args...)
-    fe = first(exprs)
-    if length(exprs) == 1 &&
-        get_column_expr(fe) === nothing &&
-        !(fe.head == :(=) || fe.head == :kw)
-
-        @warn "Returning a Table object from @by and @combine now requires `\$AsTable` on the LHS."
-
-        lhs = Expr(:$, :AsTable)
-        exprs = ((:($lhs = $fe)),)
-    end
 
     t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags) for ex in exprs)
 
@@ -1718,6 +1856,8 @@ and
 @by(df, :g, mx = mean(:x), sx = std(:x))
 ```
 
+$astable_docstring_snippet
+
 ### Examples
 
 ```julia
diff --git a/src/parsing.jl b/src/parsing.jl
index 3a250138..c83df4c2 100644
--- a/src/parsing.jl
+++ b/src/parsing.jl
@@ -91,7 +91,8 @@ is_macro_head(ex::Expr, name) = ex.head == :macrocall && ex.args[1] == Symbol(na
 
 const BYROW_SYM = Symbol("@byrow")
 const PASSMISSING_SYM = Symbol("@passmissing")
-const DEFAULT_FLAGS = (;BYROW_SYM => Ref(false), PASSMISSING_SYM => Ref(false))
+const ASTABLE_SYM = Symbol("@astable")
+const DEFAULT_FLAGS = (;BYROW_SYM => Ref(false), PASSMISSING_SYM => Ref(false), ASTABLE_SYM => Ref(false))
 
 extract_macro_flags(ex, exprflags = deepcopy(DEFAULT_FLAGS)) = (ex, exprflags)
 function extract_macro_flags(ex::Expr, exprflags = deepcopy(DEFAULT_FLAGS))
@@ -108,7 +109,6 @@ function extract_macro_flags(ex::Expr, exprflags = deepcopy(DEFAULT_FLAGS))
             return (ex, exprflags)
         end
     end
-
     return (ex, exprflags)
 end
 
@@ -126,6 +126,9 @@ function check_macro_flags_consistency(exprflags)
         if !exprflags[BYROW_SYM][]
             s = "The `@passmissing` flag is currently only allowed with the `@byrow` flag"
             throw(ArgumentError(s))
+        elseif exprflags[ASTABLE_SYM][]
+            s = "The `@passmissing` flag is currently not allowed with the `@astable` flag"
+            throw(ArgumentError(s))
         end
     end
 end
@@ -269,7 +272,13 @@ function fun_to_vec(ex::Expr;
         return ex_col
     end
 
-    if no_dest
+    if final_flags[ASTABLE_SYM][]
+        src, fun = get_source_fun_astable(ex; exprflags = final_flags)
+
+        return :($src => $fun => AsTable)
+    end
+
+    if no_dest # subset and with
         src, fun = get_source_fun(ex, exprflags = final_flags)
         return quote
             $src => $fun
@@ -359,7 +368,7 @@ function create_args_vector(arg; wrap_byrow::Bool=false)
         outer_flags[BYROW_SYM][] = true
     end
 
-    if arg isa Expr && arg.head == :block
+    if arg isa Expr && arg.head == :block && !outer_flags[ASTABLE_SYM][]
         x = MacroTools.rmlines(arg).args
     else
         x = Any[arg]
diff --git a/src/parsing_astable.jl b/src/parsing_astable.jl
new file mode 100644
index 00000000..16987d69
--- /dev/null
+++ b/src/parsing_astable.jl
@@ -0,0 +1,107 @@
+function conditionally_add_symbols!(inputs_to_function::AbstractDict,
+                                    lhs_assignments::OrderedCollections.OrderedDict, col)
+    # if it's already been assigned at top-level,
+    # don't add it to the inputs
+    if haskey(lhs_assignments, col)
+        return lhs_assignments[col]
+    else
+        return addkey!(inputs_to_function, col)
+    end
+end
+
+replace_syms_astable!(inputs_to_function::AbstractDict,
+                      lhs_assignments::OrderedCollections.OrderedDict, x) = x
+replace_syms_astable!(inputs_to_function::AbstractDict,
+                      lhs_assignments::OrderedCollections.OrderedDict, q::QuoteNode) =
+    conditionally_add_symbols!(inputs_to_function, lhs_assignments, q)
+
+function replace_syms_astable!(inputs_to_function::AbstractDict,
+                               lhs_assignments::OrderedCollections.OrderedDict, e::Expr)
+    if onearg(e, :^)
+        return e.args[2]
+    end
+
+    col = get_column_expr(e)
+    if col !== nothing
+        return conditionally_add_symbols!(inputs_to_function, lhs_assignments, col)
+    elseif e.head == :.
+        return replace_dotted_astable!(inputs_to_function, lhs_assignments, e)
+    else
+        return mapexpr(x -> replace_syms_astable!(inputs_to_function, lhs_assignments, x), e)
+    end
+end
+
+protect_replace_syms_astable!(inputs_to_function::AbstractDict,
+                              lhs_assignments::OrderedCollections.OrderedDict, e) = e
+protect_replace_syms_astable!(inputs_to_function::AbstractDict,
+                              lhs_assignments::OrderedCollections.OrderedDict, e::Expr) =
+    replace_syms!(inputs_to_function, lhs_assignments, e)
+
+function replace_dotted_astable!(inputs_to_function::AbstractDict,
+                                 lhs_assignments::OrderedCollections.OrderedDict, e)
+    x_new = replace_syms_astable!(inputs_to_function, lhs_assignments, e.args[1])
+    y_new = protect_replace_syms_astable!(inputs_to_function, lhs_assignments, e.args[2])
+    Expr(:., x_new, y_new)
+end
+
+is_column_assigment(ex) = false
+function is_column_assigment(ex::Expr)
+    ex.head == :(=) && (get_column_expr(ex.args[1]) !== nothing)
+end
+
+# Taken from MacroTools.jl
+# No docstring so assumed unstable
+block(ex) = isexpr(ex, :block) ? ex : :($ex;)
+
+sym_or_str_to_sym(x::Union{AbstractString, Symbol}) = Symbol(x)
+sym_or_str_to_sym(x) =
+    throw(ArgumentError("New columns created inside @astable must be Symbols or AbstractStrings"))
+
+function get_source_fun_astable(ex; exprflags = deepcopy(DEFAULT_FLAGS))
+    inputs_to_function = Dict{Any, Symbol}()
+    lhs_assignments = OrderedCollections.OrderedDict{Any, Symbol}()
+
+    # Make sure all top-level assignments are
+    # in the args vector
+    ex = block(MacroTools.flatten(ex))
+    exprs = map(ex.args) do arg
+        if is_column_assigment(arg)
+            lhs = get_column_expr(arg.args[1])
+            rhs = arg.args[2]
+            new_ex = replace_syms_astable!(inputs_to_function, lhs_assignments, arg.args[2])
+            if haskey(inputs_to_function, lhs)
+                new_lhs = inputs_to_function[lhs]
+                lhs_assignments[lhs] = new_lhs
+            else
+                new_lhs = addkey!(lhs_assignments, lhs)
+            end
+
+            Expr(:(=), new_lhs, new_ex)
+        else
+            replace_syms_astable!(inputs_to_function, lhs_assignments, arg)
+        end
+    end
+    source = :(DataFramesMeta.make_source_concrete($(Expr(:vect, keys(inputs_to_function)...))))
+
+    inputargs = Expr(:tuple, values(inputs_to_function)...)
+    nt_iterator = (:(DataFramesMeta.sym_or_str_to_sym($k) => $v) for (k, v) in lhs_assignments)
+    nt_expr = Expr(:tuple, Expr(:parameters, nt_iterator...))
+
+    body = Expr(:block, Expr(:block, exprs...), nt_expr)
+
+    fun = quote
+        $inputargs -> begin
+            $body
+        end
+    end
+
+    # TODO: Add passmissing support by
+    # checking if any input arguments missing,
+    # and if-so, making a named tuple with
+    # missing values
+    if exprflags[BYROW_SYM][]
+        fun = :(ByRow($fun))
+    end
+
+    return source, fun
+end
\ No newline at end of file
diff --git a/test/astable_flag.jl b/test/astable_flag.jl
new file mode 100644
index 00000000..01a2b32b
--- /dev/null
+++ b/test/astable_flag.jl
@@ -0,0 +1,216 @@
+module TestAsTableFlag
+
+using Test
+using DataFramesMeta
+using Statistics
+
+const ≅ = isequal
+
+@testset "@astable with just assignments" begin
+    df = DataFrame(a = 1, b = 2)
+
+    d = @rtransform df @astable begin
+        :x = 1
+        nothing
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 1)
+
+    d = @rselect df @astable begin
+        :x = 1
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(x = 1)
+
+    d = @transform df @astable begin
+        :x = [5]
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 5)
+
+    d = @select df @astable begin
+        :x = [5]
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(x = 5)
+end
+
+@testset "@astable with strings" begin
+    df = DataFrame(a = 1, b = 2)
+
+    x_str = "x"
+    d = @rtransform df @astable begin
+        $x_str = 1
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 1)
+
+    d = @rselect df @astable begin
+        $x_str = 1
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(x = 1)
+
+    d = @transform df @astable begin
+        $x_str = [5]
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 5)
+
+    d = @select df @astable begin
+        $x_str = [5]
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(x = 5)
+end
+
+@testset "Re-using variables" begin
+    df = DataFrame(a = 1, b = 2)
+
+    d = @rtransform df @astable begin
+        :x = 1
+        y = 5
+        :z = :x + y
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 1, z = 6)
+
+    d = @rselect df @astable begin
+        :x = 1
+        y = 5
+        :z = :x + y
+    end
+
+    @test d == DataFrame(x = 1, z = 6)
+
+    x_str = "x"
+    d = @rtransform df @astable begin
+        $x_str = 1
+        y = 5
+        :z = $x_str + y
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 1, z = 6)
+
+    d = @rselect df @astable begin
+        $x_str = 1
+        y = 5
+        :z = $x_str + y
+    end
+
+    @test d == DataFrame(x = 1, z = 6)
+end
+
+@testset "@astable with mutation" begin
+    df = DataFrame(a = 1, b = 2)
+
+    df2 = copy(df)
+    d = @rtransform! df2 @astable begin
+        :x = 1
+        nothing
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 1)
+    @test d === df2
+
+    df2 = copy(df)
+    d = @rselect! df2 @astable begin
+        :x = 1
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(x = 1)
+    @test d === df2
+
+    df2 = copy(df)
+    d = @transform! df2 @astable begin
+        :x = [5]
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(a = 1, b = 2, x = 5)
+    @test d === df2
+
+    df2 = copy(df)
+    d = @select! df2 @astable begin
+        :x = [5]
+        y = 100
+        nothing
+    end
+
+    @test d == DataFrame(x = 5)
+    @test d === df2
+end
+
+@testset "grouping astable flag" begin
+    df = DataFrame(a = [1, 1, 2, 2], b = [5, 6, 7, 8])
+
+    gd = groupby(df, :a)
+
+    d = @combine gd @astable begin
+        ex = extrema(:b)
+        :b_min = ex[1]
+        :b_max = ex[2]
+    end
+
+    res_sorted = DataFrame(a = [1, 2], b_min = [5, 7], b_max = [6, 8])
+
+    @test sort(d, :b_min) == res_sorted
+
+    d = @combine gd @astable begin
+        ex = extrema(:b)
+        $"b_min" = ex[1]
+        $"b_max" = ex[2]
+    end
+
+    @test sort(d, :b_min) == res_sorted
+
+    d = @by df :a @astable begin
+        ex = extrema(:b)
+        :b_min = ex[1]
+        :b_max = ex[2]
+    end
+
+    @test sort(d, :b_min) == res_sorted
+
+    d = @by df :a @astable begin
+        ex = extrema(:b)
+        $"b_min" = ex[1]
+        $"b_max" = ex[2]
+    end
+
+    @test sort(d, :b_min) == res_sorted
+end
+
+@testset "errors with passmissing" begin
+    @eval df = DataFrame(y = 1)
+    @test_throws LoadError @eval @transform df @passmising @byrow @astable :x = 2
+    @test_throws LoadError @eval @transform df @byrow @astable @passmissing :x = 2
+    @test_throws LoadError @eval @transform df @astable @passmissing @byrow :x = 2
+    @test_throws LoadError @eval @rtransform df @astable @passmissing :x = 2
+    @test_throws LoadError @eval @rtransform df @passmissing @astable :x = 2
+end
+
+@testset "bad assignments" begin
+    @eval df = DataFrame(y = 1)
+    @test_throws ArgumentError @eval @transform df @astable cols(1) = :y
+    @test_throws ArgumentError @eval @transform df @astable cols(AsTable) = :y
+end
+
+end # module
\ No newline at end of file
diff --git a/test/deprecated.jl b/test/deprecated.jl
index b76c8cbc..126ec441 100644
--- a/test/deprecated.jl
+++ b/test/deprecated.jl
@@ -42,7 +42,6 @@ const ≅ = isequal
     @test @based_on(gd, n = first(Symbol.(:y, ^(:body)))).n == [:vbody, :ybody]
     @test @based_on(gd, body = :i).body == df.i
     @test @based_on(gd, transform = :i).transform == df.i
-    @test @based_on(gd, (n1 = [first(:i)], n2 = [first(:y)])).n1 == [1, 4]
 
     @test @based_on(gd, n = mean(cols(iq))).n == [2.0, 4.5]
     @test @based_on(gd, n = mean(cols(iq)) + mean(cols(gq))).n == [3.0, 6.5]
@@ -51,7 +50,6 @@ const ≅ = isequal
     @test @based_on(gd, n = first(Symbol.(cols(yq), ^(:body)))).n == [:vbody, :ybody]
     @test @based_on(gd, body = cols(iq)).body == df.i
     @test @based_on(gd, transform = cols(iq)).transform == df.i
-    @test @based_on(gd, (n1 = [first(cols(iq))], n2 = [first(cols(yq))])).n1 == [1, 4]
 
     @test @based_on(gd, n = mean(cols(ir))).n == [2.0, 4.5]
     @test @based_on(gd, n = mean(cols(ir)) + mean(cols(gr))).n == [3.0, 6.5]
@@ -60,7 +58,6 @@ const ≅ = isequal
     @test @based_on(gd, n = first(Symbol.(cols(yr), ^(:body)))).n == [:vbody, :ybody]
     @test @based_on(gd, body = cols(ir)).body == df.i
     @test @based_on(gd, transform = cols(ir)).transform == df.i
-    @test @based_on(gd, (n1 = [first(cols(ir))], n2 = [first(cols(yr))])).n1 == [1, 4]
     @test @based_on(gd, n = mean(cols("i")) + 0 * first(cols(:g))).n == [2.0, 4.5]
     @test @based_on(gd, n = mean(cols(2)) + first(cols(1))).n == [3.0, 6.5]
 
diff --git a/test/function_compilation.jl b/test/function_compilation.jl
index 4c411f61..5921942d 100644
--- a/test/function_compilation.jl
+++ b/test/function_compilation.jl
@@ -154,9 +154,9 @@ using DataFramesMeta
 
             gd = groupby(df, :a)
 
-            @test @combine(gd, testnt(:b)) == DataFrame(a = [1], c =  [2])
+            @test @combine(gd, cols(AsTable) = testnt(:b)) == DataFrame(a = [1], c =  [2])
 
-            fasttime = @timed @combine(gd, testnt(:b))
+            fasttime = @timed @combine(gd, cols(AsTable) = testnt(:b))
             slowtime = @timed combine(gd, :b => (b -> testnt(b)) => AsTable)
             (slowtime[2] > fasttime[2]) || @warn("Slow compilation")
 
diff --git a/test/grouping.jl b/test/grouping.jl
index a998c8a0..e4ea60b0 100644
--- a/test/grouping.jl
+++ b/test/grouping.jl
@@ -49,7 +49,6 @@ g = groupby(d, :x, sort=true)
     @test @combine(gd, :n = first(Symbol.(:y, ^(:body)))).n == [:vbody, :ybody]
     @test @combine(gd, :body =  :i).body == df.i
     @test @combine(gd, :transform = :i).transform == df.i
-    @test @combine(gd, (n1 = [first(:i)], n2 = [first(:y)])).n1 == [1, 4]
 
     @test @combine(gd, :n = mean($iq)).n == [2.0, 4.5]
     @test @combine(gd, :n = mean($iq) + mean($gq)).n == [3.0, 6.5]
@@ -59,7 +58,6 @@ g = groupby(d, :x, sort=true)
     @test @combine(gd, $:n = mean($:i)).n == [2.0, 4.5]
     @test @combine(gd, :body =  $iq).body == df.i
     @test @combine(gd, :transform = $iq).transform == df.i
-    @test @combine(gd, (n1 = [first($iq)], n2 = [first($yq)])).n1 == [1, 4]
 
     @test @combine(gd, :n = mean($ir)).n == [2.0, 4.5]
     @test @combine(gd, :n = mean($ir) + mean($gr)).n == [3.0, 6.5]
@@ -68,7 +66,6 @@ g = groupby(d, :x, sort=true)
     @test @combine(gd, :n = first(Symbol.($yr, ^(:body)))).n == [:vbody, :ybody]
     @test @combine(gd, :body =  $ir).body == df.i
     @test @combine(gd, :transform = $ir).transform == df.i
-    @test @combine(gd, (n1 = [first($ir)], n2 = [first($yr)])).n1 == [1, 4]
     @test @combine(gd, :n = mean($"i") + 0 * first($:g)).n == [2.0, 4.5]
     @test @combine(gd, :n = mean($2) + first($1)).n == [3.0, 6.5]
 
@@ -192,7 +189,6 @@ end
     @test @by(df, :g, :n = first(Symbol.(:y, ^(:body)))).n == [:vbody, :ybody]
     @test @by(df, :g, :body =  :i).body == df.i
     @test @by(df, :g, :transform = :i).transform == df.i
-    @test @by(df, :g, (n1 = [first(:i)], n2 = [first(:y)])).n1 == [1, 4]
 
     @test @by(df, :g, :n = mean($iq)).n == [2.0, 4.5]
     @test @by(df, :g, :n = mean($iq) + mean($gq)).n == [3.0, 6.5]
@@ -202,7 +198,6 @@ end
     @test @by(df, :g, $:n = mean($:i)).n == [2.0, 4.5]
     @test @by(df, :g, :body =  $iq).body == df.i
     @test @by(df, :g, :transform = $iq).transform == df.i
-    @test @by(df, :g, (n1 = [first($iq)], n2 = [first($yq)])).n1 == [1, 4]
 
     @test @by(df, "g", :n = mean($ir)).n == [2.0, 4.5]
     @test @by(df, "g", :n = mean($ir) + mean($gr)).n == [3.0, 6.5]
@@ -211,7 +206,6 @@ end
     @test @by(df, "g", :n = first(Symbol.($yr, ^(:body)))).n == [:vbody, :ybody]
     @test @by(df, "g", :body =  $ir).body == df.i
     @test @by(df, "g", :transform = $ir).transform == df.i
-    @test @by(df, "g", (n1 = [first($ir)], n2 = [first($yr)])).n1 == [1, 4]
     @test @by(df, "g", :n = mean($"i") + 0 * first($:g)).n == [2.0, 4.5]
     @test @by(df, "g", :n = mean($2) + first($1)).n == [3.0, 6.5]
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 3218556f..5ab6f363 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -12,6 +12,7 @@ my_tests = ["dataframes.jl",
             "deprecated.jl",
             "byrow.jl",
             "astable.jl",
+            "astable_flag.jl",
             "passmissing.jl"]
 
 println("Running tests:")