Update

src/macros/@nonlinear.jl → src/macros/@force_nonlinear.jl

@@ -7,15 +7,20 @@ const _op_add = NonlinearOperator(+, :+)
 const _op_sub = NonlinearOperator(-, :-)
 const _op_mul = NonlinearOperator(*, :*)
 const _op_div = NonlinearOperator(/, :/)
+const _op_pow = NonlinearOperator(^, :^)
 
 """
-    @nonlinear(expr)
+    @force_nonlinear(expr)
 
 Change the parsing of `expr` to construct [`GenericNonlinearExpr`](@ref) instead
 of [`GenericAffExpr`](@ref) or [`GenericQuadExpr`](@ref).
 
-This macro works by walking `expr` and substituting all calls to `+`, `-`, `*`
-and `/` in favor of ones that construct [`GenericNonlinearExpr`](@ref).
+This macro works by walking `expr` and substituting all calls to `+`, `-`, `*`,
+`/`, and `^` in favor of ones that construct [`GenericNonlinearExpr`](@ref).
+
+This macro will error if the resulting expression does not produce a
+[`GenericNonlinearExpr`](@ref) because, for example, it is used on an expression
+that does not use the basic arithmetic operators.
 
 ## When to use this macro
 
@@ -38,13 +43,13 @@ julia> @variable(model, x);
 julia> @expression(model, (x - 0.1)^2)
 x² - 0.2 x + 0.010000000000000002
 
-julia> @expression(model, @nonlinear((x - 0.1)^2))
+julia> @expression(model, @force_nonlinear((x - 0.1)^2))
 (x - 0.1) ^ 2.0
 
 julia> (x - 0.1)^2
 x² - 0.2 x + 0.010000000000000002
 
-julia> @nonlinear((x - 0.1)^2)
+julia> @force_nonlinear((x - 0.1)^2)
 (x - 0.1) ^ 2.0
 ```
 
@@ -75,17 +80,18 @@ julia> @variable(model, x);
 julia> @expression(model, x * 2.0 * (1 + x) * x)
 (2 x² + 2 x) * x
 
-julia> @expression(model, @nonlinear(x * 2.0 * (1 + x) * x))
+julia> @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x))
 x * 2.0 * (1 + x) * x
 
 julia> @allocated @expression(model, x * 2.0 * (1 + x) * x)
 3200
 
-julia> @allocated @expression(model, @nonlinear(x * 2.0 * (1 + x) * x))
+julia> @allocated @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x))
 640
 ```
 """
-macro nonlinear(expr)
+macro force_nonlinear(expr)
+    error_fn = Containers.build_error_fn(:force_nonlinear, (expr,), __source__)
     ret = MacroTools.postwalk(expr) do x
         if Meta.isexpr(x, :call)
             if x.args[1] == :+
@@ -96,9 +102,17 @@ macro nonlinear(expr)
                 return Expr(:call, _op_mul, x.args[2:end]...)
             elseif x.args[1] == :/
                 return Expr(:call, _op_div, x.args[2:end]...)
+            elseif x.args[1] == :^
+                return Expr(:call, _op_pow, x.args[2:end]...)
             end
         end
         return x
     end
-    return esc(ret)
+    return quote
+        r = $(esc(ret))
+        if !(r isa $GenericNonlinearExpr)
+            $error_fn("expression did not produce a GenericNonlinearExpr")
+        end
+        r
+    end
 end

test/test_macros.jl

@@ -2358,4 +2358,26 @@ function test_op_or_short_circuit()
     return
 end
 
+function test_force_nonlinear()
+    model = Model()
+    @variable(model, x)
+    @test (1 + x) isa AffExpr
+    @test @force_nonlinear(1 + x) isa GenericNonlinearExpr
+    @test (1 - x) isa AffExpr
+    @test @force_nonlinear(1 - x) isa GenericNonlinearExpr
+    @test (2 * x) isa AffExpr
+    @test @force_nonlinear(2 * x) isa GenericNonlinearExpr
+    @test (x / 3) isa AffExpr
+    @test @force_nonlinear(x / 3) isa GenericNonlinearExpr
+    @test (x ^ 2) isa QuadExpr
+    @test @force_nonlinear(x ^ 2) isa GenericNonlinearExpr
+    @test_throws_runtime(
+        ErrorException(
+            "In `@force_nonlinear(x)`: expression did not produce a GenericNonlinearExpr",
+        ),
+        @force_nonlinear(x),
+    )
+    return
+end
+
 end  # module