Make @PolyVar return a tuple of all created variables

JuliaAlgebra · Aug 7, 2018 · e8f8975 · e8f8975
1 parent e08321a
commit e8f8975
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/src/var.jl b/src/var.jl
@@ -11,7 +11,7 @@ end
 
 function buildpolyvar(::Type{PV}, var) where {PV}
     if isa(var, Symbol)
-        :($(esc(var)) = $PV($"$var"))
+        var, :($(esc(var)) = $PV($"$var"))
     else
         isa(var, Expr) || error("Expected $var to be a variable name")
         Base.Meta.isexpr(var, :ref) || error("Expected $var to be of the form varname[idxset]")
@@ -20,21 +20,35 @@ function buildpolyvar(::Type{PV}, var) where {PV}
         prefix = string(var.args[1])
         if length(var.args) == 2
             idxset = esc(var.args[2])
-            :($(esc(varname)) = polyvecvar($PV, $prefix, $idxset))
+            varname, :($(esc(varname)) = polyvecvar($PV, $prefix, $idxset))
         else
             rowidxset = esc(var.args[2])
             colidxset = esc(var.args[3])
-            :($(esc(varname)) = polymatrixvar($PV, $prefix, $rowidxset, $colidxset))
+            varname, :($(esc(varname)) = polymatrixvar($PV, $prefix, $rowidxset, $colidxset))
         end
     end
 end
 
+function buildpolyvars(::Type{PV}, args) where {PV}
+    vars = Symbol[]
+    exprs = []
+    for arg in args
+        var, expr = buildpolyvar(PV, arg)
+        push!(vars, var)
+        push!(exprs, expr)
+    end
+    vars, exprs
+end
+
 # Variable vector x returned garanteed to be sorted so that if p is built with x then vars(p) == x
 macro polyvar(args...)
-    reduce((x,y) -> :($x; $y), [buildpolyvar(PolyVar{true}, arg) for arg in args], init=:())
+    vars, exprs = buildpolyvars(PolyVar{true}, args)
+    :($(foldl((x,y) -> :($x; $y), exprs, init=:())); $(Expr(:tuple, esc.(vars)...)))
 end
+
 macro ncpolyvar(args...)
-    reduce((x,y) -> :($x; $y), [buildpolyvar(PolyVar{false}, arg) for arg in args], init=:())
+    vars, exprs = buildpolyvars(PolyVar{false}, args)
+    :($(foldl((x,y) -> :($x; $y), exprs, init=:())); $(Expr(:tuple, esc.(vars)...)))
 end
 
 struct PolyVar{C} <: AbstractVariable

diff --git a/test/mono.jl b/test/mono.jl
@@ -1,5 +1,5 @@
 @testset "PolyVar and Monomial tests" begin
-    @testset "polyvar macro index set" begin
+    @testset "PolyVar macro index set" begin
         n = 3
         @polyvar x[1:n] y z[1:n-1] u[1:n,1:n-1]
         @test x isa Vector{PolyVar{true}}
@@ -12,6 +12,16 @@
         @test x[1] > x[2] > x[3] > y > z[1] > z[2]
         @test u[1, 1] > u[2, 1] > u[2, 2]
     end
+    @testset "PolyVar macro tuple return" begin
+        vars = @polyvar x y z
+        @test vars isa Tuple
+        @test vars == (x, y, z)
+
+        vars = @ncpolyvar x y z
+        @test vars isa Tuple
+        @test vars == (x, y, z)
+    end
+
     @testset "PolyVar" begin
         @test zeroterm(PolyVar{true}) == 0
         @test zero(PolyVar{true}) == 0