Another attempt at an astable flag

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 2nd 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])

src/macros.jl

@@ -284,7 +284,7 @@ end
 
 
 """
-    passmissing(args...)
+    @passmissing(args...)
 
 Propograte missing values inside DataFramesMeta.jl macros.
 
@@ -351,13 +351,14 @@ macro passmissing(args...)
 end
 
 """
-    astable(args...)
+    @astable(args...)
 
-Return a `NamedTuple` from a transformation inside DataFramesMeta.jl macros.
+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 = x`, i.e. assignments to a
-`Symbol`, which is a syntax error outside of the macro. At the end of the
+and collects all such expressions of the form `:y = ...`, i.e. assignments to a
+`Symbol`, 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.
@@ -374,7 +375,7 @@ df = DataFrame(a = 1)
 end
 ```
 
-becomes the pair
+become the pair
 
 ```
 function f(a)
@@ -388,8 +389,24 @@ end
 transform(df, [:a] => ByRow(f) => AsTable)
 ```
 
-`@astable` is useful when performing intermediate calculations
-yet store their results in new columns. For example, the following fails.
+`@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
@@ -406,6 +423,13 @@ transformations. `@astable` solves this problem
     :new_col_2 = :new_col_1 + :z
 end
 
+Column assignment in `@astable` follows the same rules as
+column assignment more generally. Construct a new column
+from a string by escaping it with `$DOLLAR`, which can be a
+`Symbol` or an `AbstractString`. References to existing
+columns may be a `Symbol`, `AbstractString`, or an
+integer.
+
 ### Examples
 
 ```
@@ -427,15 +451,30 @@ julia> d = @rtransform df @astable begin
 julia> df = DataFrame(a = [1, 1, 2, 2], b = [5, 6, 70, 80]);
 
 julia> @by df :a @astable begin
-           $(DOLLAR)"Mean of b" = mean(:b)
-           $(DOLLAR)"Standard deviation of b" = std(:b)
+            ex = extrema(:b)
+            :min_b = first(ex)
+            :max_b = last(ex)
        end
 2×3 DataFrame
- Row │ a      Mean of b  Standard deviation of b
-     │ Int64  Float64    Float64
-─────┼───────────────────────────────────────────
-   1 │     1        5.5                 0.707107
-   2 │     2       75.0                 7.07107
+ Row │ a      min_b  max_b
+     │ Int64  Int64  Int64
+─────┼─────────────────────
+   1 │     1      5      6
+   2 │     2     70     80
+
+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
 ```
 
 """

src/parsing.jl

@@ -276,7 +276,7 @@ function fun_to_vec(ex::Expr;
         return :($src => $fun => AsTable)
     end
 
-    if no_dest # subet and with
+    if no_dest # subset and with
         src, fun = get_source_fun(ex, exprflags = final_flags)
         return quote
             $src => $fun

src/parsing_astable.jl

@@ -1,4 +1,5 @@
-function conditionally_add_symbols!(inputs_to_function, lhs_assignments, col)
+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)
@@ -8,11 +9,14 @@ function conditionally_add_symbols!(inputs_to_function, lhs_assignments, col)
     end
 end
 
-replace_syms_astable!(inputs_to_function, lhs_assignments, x) = x
-replace_syms_astable!(inputs_to_function, lhs_assignments, q::QuoteNode) =
+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, lhs_assignments, e::Expr)
+function replace_syms_astable!(inputs_to_function::AbstractDict,
+                               lhs_assignments::OrderedCollections.OrderedDict, e::Expr)
     if onearg(e, :^)
         return e.args[2]
     end
@@ -27,11 +31,14 @@ function replace_syms_astable!(inputs_to_function, lhs_assignments, e::Expr)
     end
 end
 
-protect_replace_syms_astable!(inputs_to_function, lhs_assignments, e) = e
-protect_replace_syms_astable!(inputs_to_function, lhs_assignments, e::Expr) =
+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, 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)
@@ -43,9 +50,15 @@ function is_column_assigment(ex::Expr)
 end
 
 # Taken from MacroTools.jl
-# No docstring so assumed untable
+# No docstring so assumed unstable
 block(ex) = isexpr(ex, :block) ? ex : :($ex;)
 
+sym_or_str_to_sym(x::Union{AbstractString, Symbol}) = Symbol(x)
+function sym_or_str_to_sym(x)
+    e = "New columns created inside @astable must be Symbols or AbstractStrings"
+    throw(ArgumentError(e))
+end
+
 function get_source_fun_astable(ex; exprflags = deepcopy(DEFAULT_FLAGS))
     inputs_to_function = Dict{Any, Symbol}()
     lhs_assignments = OrderedCollections.OrderedDict{Any, Symbol}()
@@ -73,7 +86,7 @@ function get_source_fun_astable(ex; exprflags = deepcopy(DEFAULT_FLAGS))
     source = :(DataFramesMeta.make_source_concrete($(Expr(:vect, keys(inputs_to_function)...))))
 
     inputargs = Expr(:tuple, values(inputs_to_function)...)
-    nt_iterator = (:(Symbol($k) => $v) for (k, v) in lhs_assignments)
+    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)
 

test/astable_flag.jl

@@ -120,5 +120,50 @@ end
     @test d == DataFrame(x = 1, z = 6)
 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
+
+    @test sort(d.b_min) == [5, 7]
+
+    d = @combine gd @astable begin
+        ex = extrema(:b)
+        $"b_min" = ex[1]
+        $"b_max" = ex[2]
+    end
+
+    @test sort(d.b_min) == [5, 7]
+
+    d = @by df :a @astable begin
+        ex = extrema(:b)
+        :b_min = ex[1]
+        :b_max = ex[2]
+    end
+
+    @test sort(d.b_min) == [5, 7]
+
+    d = @by df :a @astable begin
+        ex = extrema(:b)
+        $"b_min" = ex[1]
+        $"b_max" = ex[2]
+    end
+
+    @test sort(d.b_min) == [5, 7]
+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