From 39a193fc692b0f0ddeacfcc4447352c8266e1db8 Mon Sep 17 00:00:00 2001
From: Alexander Plavin <alexander@plav.in>
Date: Thu, 30 May 2024 15:56:36 -0400
Subject: [PATCH] generalize lift() and @lift()

@lift: allow single observable, or no observables at all
lift: allow operating on non-observables
---
 src/interaction/liftmacro.jl   | 21 +++++----------------
 src/interaction/observables.jl | 14 ++++++++++++--
 src/scenes.jl                  |  4 ++++
 test/observables.jl            | 10 ++++++++++
 4 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/src/interaction/liftmacro.jl b/src/interaction/liftmacro.jl
index eaf867127d2..850874db054 100644
--- a/src/interaction/liftmacro.jl
+++ b/src/interaction/liftmacro.jl
@@ -24,22 +24,15 @@ end
 """
 Replaces every subexpression that looks like a observable expression with a substitute symbol stored in `exprdict`.
 """
-function replace_observable_expressions!(exp::Expr, exprdict)
+function replace_observable_expressions(exp::Expr, exprdict)
     if is_interpolated_observable(exp)
-        error("You can't @lift an expression that only consists of a single observable.")
+        exprdict[exp]
     else
-        for (i, arg) in enumerate(exp.args)
-            if is_interpolated_observable(arg)
-                exp.args[i] = exprdict[arg]
-            else
-                replace_observable_expressions!(arg, exprdict)
-            end
-        end
+        Expr(exp.head, replace_observable_expressions.(exp.args, Ref(exprdict))...)
     end
-    return exp
 end
 
-replace_observable_expressions!(x, exprdict) = nothing
+replace_observable_expressions(x, exprdict) = x
 
 """
 Replaces an expression with `lift(argtuple -> expression, args...)`, where `args`
@@ -73,15 +66,11 @@ macro lift(exp)
 
     observable_expr_set = find_observable_expressions(exp)
 
-    if length(observable_expr_set) == 0
-        error("Did not find any interpolated observables. Use '\$(observable)' to interpolate it into the macro.")
-    end
-
     # store expressions with their substitute symbols, gensym them manually to be
     # able to escape the expression later
     observable_expr_arg_dict = Dict(expr => gensym("arg$i") for (i, expr) in enumerate(observable_expr_set))
 
-    replace_observable_expressions!(exp, observable_expr_arg_dict)
+    exp = replace_observable_expressions(exp, observable_expr_arg_dict)
 
     # keep an array for ordering
     observable_expressions_array = collect(keys(observable_expr_arg_dict))
diff --git a/src/interaction/observables.jl b/src/interaction/observables.jl
index 203a9257a4d..9551d542cfa 100644
--- a/src/interaction/observables.jl
+++ b/src/interaction/observables.jl
@@ -1,5 +1,15 @@
-# lift makes it easier to search + replace observable code, while `map` is really hard to differentiate from `map(f, array)`
-const lift = map
+# mimick Observables.jl map() signature to forward directly:
+lift(f, arg::AbstractObservable, args...; kwargs...) = map(f, arg, args...; kwargs...)
+# handle the general case:
+function lift(f, args...; kwargs...)
+    if !any(a -> isa(a, AbstractObservable), args)
+        # there are no observables
+        f(args...)
+    else
+        # there are observables, but not in the first position
+        lift((_, as...) -> f(as...), Observable(nothing), args...; kwargs...)
+    end
+end
 
 """
 Observables.off but without throwing an error
diff --git a/src/scenes.jl b/src/scenes.jl
index 4b139cceb3b..1b6cc302a9b 100644
--- a/src/scenes.jl
+++ b/src/scenes.jl
@@ -150,6 +150,10 @@ function Observables.onany(@nospecialize(f), @nospecialize(scene::Union{Plot,Sce
     return to_deregister
 end
 
+lift(f, @nospecialize(arg::Union{Plot,Scene}), args...; kwargs...) = map(f, arg, args...; kwargs...)
+
+# map! and map are for backward compatibility
+# can move everything to use lift() in the future
 @inline function Base.map!(f, @nospecialize(scene::Union{Plot,Scene}), result::AbstractObservable, os...;
                            update::Bool=true, priority = 0)
     # note: the @inline prevents de-specialization due to the splatting
diff --git a/test/observables.jl b/test/observables.jl
index 386720c99c1..baed067e8b6 100644
--- a/test/observables.jl
+++ b/test/observables.jl
@@ -1,8 +1,18 @@
 @testset "lift macro" begin
+    u_noobs = "a"
     x = Observable(1.0)
     y = Observable(2.0)
     z = (x = x, y = y)
 
+    noobs = @lift u_noobs * "b"
+    @test noobs == "ab"
+
+    noobs = @lift $u_noobs * "b"
+    @test noobs == "ab"
+
+    xx = @lift $x
+    @test xx[] == 1.0
+
     t1 = @lift($x + $y)
     @test t1[] == 3.0
     t2 = @lift($(z.x) - $(z.y))