add at_scoped macro

Project.toml

@@ -1,7 +1,7 @@
 name = "ScopedValues"
 uuid = "7e506255-f358-4e82-b7e4-beb19740aa63"
 authors = ["Valentin Churavy <[email protected]>"]
-version = "1.0.1"
+version = "1.1.0"
 
 [deps]
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"

src/ScopedValues.jl

@@ -1,9 +1,9 @@
 module ScopedValues
 
-export ScopedValue, scoped
+export ScopedValue, scoped, @scoped
 
 if isdefined(Base, :ScopedValues)
-    import Base.ScopedValues: ScopedValue, scoped, Scope, current_scope
+    import Base.ScopedValues: ScopedValue, scoped, @scoped, Scope, current_scope
 else
 
 """
@@ -48,6 +48,13 @@ end
 Scope(parent, key::ScopedValue{T}, value) where T =
     Scope(parent, key, convert(T, value))
 
+function Scope(scope, pairs::Pair{<:ScopedValue}...)
+    for pair in pairs
+        scope = Scope(scope, pair...)
+    end
+    return scope
+end
+
 """
     current_scope()::Union{Nothing, Scope}
 
@@ -99,11 +106,8 @@ end
 Execute `f` in a new scope with `var` set to `val`.
 """
 function scoped(f, pair::Pair{<:ScopedValue}, rest::Pair{<:ScopedValue}...)
-    scope = current_scope()
-    scope = Scope(scope, pair...)
-    for pair in rest
-        scope = Scope(scope, pair...)
-    end
+    @nospecialize
+    scope = Scope(current_scope(), pair, rest...)
     enter_scope(scope) do
         f()
     end
@@ -113,6 +117,43 @@ scoped(f) = f()
 
 include("payloadlogger.jl")
 
+
+import Base.CoreLogging: LogState
+"""
+    @scoped vars... expr
+
+Macro version of `scoped(f, vars...)` but with `expr` instead of `f` function.
+This is similar to using [`scoped`](@ref) with a `do` block, but avoids creating
+a closure.
+"""
+macro scoped(exprs...)
+    ex = last(exprs)
+    if length(exprs) > 1
+        exprs = exprs[1:end-1]
+    else
+        exprs = ()
+    end
+    for expr in exprs
+        if expr.head !== :call || first(expr.args) !== :(=>)
+            error("@scoped expects arguments of the form `A => 2` got $expr")
+        end
+    end
+    exprs = map(esc, exprs)
+    ct = gensym(:ct)
+    current_logstate = gensym(:current_logstate)
+    body = Expr(:tryfinally, esc(ex), :($(ct).logstate = $(current_logstate)))
+    quote
+        $(ct) = $(Base.current_task)()
+        $(current_logstate) = $(ct).logstate::Union{Nothing, LogState}
+        scope = $(Scope)($(current_scope)(), $(exprs...))
+        if scope !== nothing
+            logger = $(ScopePayloadLogger)($(current_logger)(), scope)
+            $(ct).logstate = $(LogState)(logger)
+        end
+        $body
+    end
+end
+
 end # isdefined
 
 end # module ScopedValues

test/runtests.jl

@@ -107,3 +107,16 @@ end
         end
     end
 end
+
+@testset "macro" begin
+    @scoped svar=>2 svar_float=>2.0 begin
+        @test svar[] == 2
+        @test svar_float[] == 2.0
+    end
+    # Doesn't do much...
+    @scoped begin
+        @test svar[] == 1
+        @test svar_float[] == 1.0
+    end
+end
+