Add @unpack_fields macro

devmotion · devmotion · commit 2d5e9630072c · 2023-03-05T12:24:32.000+01:00
diff --git a/README.md b/README.md
@@ -4,8 +4,8 @@
 [![Coverage](https://codecov.io/gh/devmotion/SimpleUnPack.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/devmotion/SimpleUnPack.jl)
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)
 
-This package provides the `@unpack` macro for destructuring properties.
-Its behaviour is equivalent to the destructuring that was introduced in [Julia#39285](https://github.com/JuliaLang/julia/pull/39285) and is available in Julia >= 1.7.0-DEV.364.
+This package provides the `@unpack` and `@unpack_fields` macros for destructuring properties and fields, respectively.
+The behaviour of `@unpack` is equivalent to the destructuring that was introduced in [Julia#39285](https://github.com/JuliaLang/julia/pull/39285) and is available in Julia >= 1.7.0-DEV.364.
 
 ## Examples
 
@@ -24,18 +24,28 @@ julia> a
 
 julia> b
 42
+
+julia> @unpack_fields a, b = f(10)
+(b = 10, a = 5.0)
+
+julia> a
+5.0
+
+julia> b
+10
 ```
 
 An example with a custom struct in a function:
 
 ```julia
 julia> using SimpleUnPack
 
-julia> struct MyStruct{T}
-           x::T
+julia> struct MyStruct{X,Y}
+           x::X
+           y::Y
        end
 
-julia> Base.getpropertynames(::MyStruct) = (:x, :y)
+julia> Base.propertynames(::MyStruct) = (:x, :y)
 
 julia> function Base.getproperty(m::MyStruct, p::Symbol)
            if p === :y
@@ -50,8 +60,16 @@ julia> function g(m::MyStruct)
            return (; x, y)
        end;
 
-julia> g(MyStruct(1.0))
+julia> g(MyStruct(1.0, -5))
 (x = 1.0, y = 42)
+
+julia> function h(m::MyStruct)
+           @unpack_fields x, y = m
+           return (; x, y)
+       end
+
+julia> h(MyStruct(1.0, -5))
+(x = 1.0, y = -5)
 ```
 
 ## Comparison with UnPack.jl
@@ -60,4 +78,7 @@ The syntax of `@unpack` is based on [`UnPack.@unpack`](https://github.com/mauro3
 However, `UnPack.@unpack` is more flexible and based on `UnPack.unpack` instad of `getproperty`.
 While `UnPack.unpack` falls back to `getproperty`, it also supports `AbstractDict`s with keys of type `Symbol` and `AbstractString`, and can be specialized for other types.
 Since `UnPack.unpack` dispatches on `Val(property)` instances, this increased flexibility comes at the cost of increased compilation times.
-Moreover, UnPack also provides an `UnPack.@pack!` macro for setting properties.
+
+In contrast to SimpleUnPack, UnPack provides an `UnPack.@pack!` macro for setting properties.
+
+However, currently UnPack does not support destructuring based on `getfield` only ([UnPack#23](https://github.com/mauro3/UnPack.jl/issues/23)).
diff --git a/src/SimpleUnPack.jl b/src/SimpleUnPack.jl
@@ -1,29 +1,68 @@
 module SimpleUnPack
 
-export @unpack
+export @unpack, @unpack_fields
 
 """
     @unpack a, b, ... = x
 
 Destructure properties `a`, `b`, ... of `x` into variables of the same name.
 
 The behaviour of the macro is equivalent to `(; a, b, ...) = x` which was introduced in [Julia#39285](https://github.com/JuliaLang/julia/pull/39285) and is available in Julia >= 1.7.0-DEV.364.
+
+See also [`@unpack_fields`](@ref)
 """
 macro unpack(args)
-    return unpack(args)
+    # Extract names of properties and RHS
+    names, rhs = split_names_rhs(:unpack, args)
+
+    # Construct destructuring expression
+    expr = if VERSION >= v"1.7.0-DEV.364"
+        # Fall back to destructuring in Base when available:
+        # https://github.com/JuliaLang/julia/pull/39285
+        Expr(:(=), Expr(:tuple, Expr(:parameters, (esc(p) for p in names)...)), esc(rhs))
+    else
+        destructuring_expr(:getproperty, names, rhs)
+    end
+
+    return expr
+end
+
+"""
+    @unpack_fields a, b, ... = x
+
+Destructure fields `a`, `b`, ... of `x` into variables of the same name.
+
+See also [`@unpack`](@ref)
+"""
+macro unpack_fields(args)
+    # Extract names of fields and RHS
+    names, rhs = split_names_rhs(:unpack_fields, args)
+
+    # Construct destructuring expression
+    expr = destructuring_expr(:getfield, names, rhs)
+
+    return expr
 end
 
-function unpack(args)
-    # Extract properties and RHS
-    if !Meta.isexpr(args, :(=), 2)
+"""
+    split_names_rhs(macrosym::Symbol, expr)
+
+Split an expression `expr` of the form `a, b, ... = x` into a tuple consisting of a vector of symbols `a`, `b`, ..., and the right-hand side `x`.
+
+The symbol `macro_name` specifies the macro from which this function is called.
+
+This function is used internally with `macrosym = :unpack` and `macrosym = :unpack_fields`.
+"""
+function split_names_rhs(macrosym::Symbol, expr)
+    if !Meta.isexpr(expr, :(=), 2)
         throw(
             ArgumentError(
-                "`@unpack` can only be applied to expressions of the form `a, b, ... = x`"
+                "`@$macrosym` can only be applied to expressions of the form `a, b, ... = x`",
             ),
         )
     end
-    lhs, rhs = args.args
-    properties = if lhs isa Symbol
+    lhs, rhs = expr.args
+    names = if lhs isa Symbol
         [lhs]
     elseif Meta.isexpr(lhs, :tuple) &&
         !isempty(lhs.args) &&
@@ -32,34 +71,33 @@ function unpack(args)
     else
         throw(
             ArgumentError(
-                "`@unpack` can only be applied to expressions of the form `a, b, ... = x`",
+                "`@$macrosym` can only be applied to expressions of the form `a, b, ... = x`",
             ),
         )
     end
+    return names, rhs
+end
 
-    if VERSION >= v"1.7.0-DEV.364"
-        # Fall back to destructuring in Base when available:
-        # https://github.com/JuliaLang/julia/pull/39285
-        return Expr(
-            :(=), Expr(:tuple, Expr(:parameters, (esc(p) for p in properties)...)), esc(rhs)
-        )
-    else
-        @gensym object
-        block = Expr(:block)
-        for p in properties
-            push!(
-                block.args,
-                Expr(:(=), esc(p), Expr(:call, :getproperty, esc(object), QuoteNode(p))),
-            )
-        end
-        return Base.remove_linenums!(
-            quote
-                $(esc(object)) = $(esc(rhs)) # In case the RHS is an expression
-                $block
-                $(esc(object)) # Return evaluation of rhs to ensure the behaviour is the same as (; ...) = rhs
-            end,
-        )
+"""
+    destructuring_expr(fsym::Symbol, names, rhs)
+
+Return an expression that destructures `rhs` based on a function of name `fsym` and keys `names` into variables of the same `names`.
+
+This function is used internally with `fsym = :getproperty` and `fsym = :getfield`.
+"""
+function destructuring_expr(fsym::Symbol, names, rhs)
+    @gensym object
+    block = Expr(:block)
+    for p in names
+        push!(block.args, Expr(:(=), esc(p), Expr(:call, fsym, esc(object), QuoteNode(p))))
     end
+    return Base.remove_linenums!(
+        quote
+            $(esc(object)) = $(esc(rhs)) # In case the RHS is an expression
+            $block
+            $(esc(object)) # Return evaluation of the RHS
+        end,
+    )
 end
 
-end
+end # module
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,9 +1,10 @@
 using SimpleUnPack
 using Test
 
-struct Property{X,Y}
+struct Property{X,Y,Z}
     x::X
     y::Y
+    z::Z
 end
 
 Base.propertynames(::Property) = (:x, :y, :z)
@@ -56,19 +57,40 @@ end
         @unpack y = d
         @test y == 1.0
 
-        d = Struct(43, 2.0, "z2")
-        @unpack x, z = d
+        d = (x=43, y=2.0, z="z2")
+        @unpack_fields x, z = d
         @test x == 43
         @test z == "z2"
-        @unpack y = d
+        @unpack_fields y = d
         @test y == 2.0
 
-        d = Property(44, 3.0)
+        d = Struct(44, 3.0, "z3")
         @unpack x, z = d
         @test x == 44
-        @test z == "z"
+        @test z == "z3"
         @unpack y = d
         @test y == 3.0
+
+        d = Struct(45, 4.0, "z4")
+        @unpack_fields x, z = d
+        @test x == 45
+        @test z == "z4"
+        @unpack_fields y = d
+        @test y == 4.0
+
+        d = Property(46, 5.0, "z5")
+        @unpack x, z = d
+        @test x == 46
+        @test z == "z"
+        @unpack y = d
+        @test y == 5.0
+
+        d = Property(47, 6.0, "z6")
+        @unpack_fields x, z = d
+        @test x == 47
+        @test z == "z6"
+        @unpack_fields y = d
+        @test y == 6.0
     end
 
     @testset "Expression as RHS" begin
@@ -78,27 +100,53 @@ end
         @unpack y = (x=42, y=1.0, z="z1")
         @test y == 1.0
 
-        @unpack x, z = Struct(43, 2.0, "z2")
+        @unpack_fields x, z = (x=43, y=2.0, z="z2")
         @test x == 43
         @test z == "z2"
-        @unpack y = Struct(43, 2.0, "z2")
+        @unpack_fields y = (x=43, y=2.0, z="z2")
         @test y == 2.0
 
-        @unpack x, z = Property(44, 3.0)
+        @unpack x, z = Struct(44, 3.0, "z3")
         @test x == 44
-        @test z == "z"
-        @unpack y = Property(44, 3.0)
+        @test z == "z3"
+        @unpack y = Struct(44, 3.0, "z3")
         @test y == 3.0
+
+        @unpack_fields x, z = Struct(45, 4.0, "z4")
+        @test x == 45
+        @test z == "z4"
+        @unpack_fields y = Struct(45, 4.0, "z4")
+        @test y == 4.0
+
+        @unpack x, z = Property(46, 5.0, "z5")
+        @test x == 46
+        @test z == "z"
+        @unpack y = Property(46, 5.0, "z5")
+        @test y == 5.0
+
+        @unpack_fields x, z = Property(47, 6.0, "z6")
+        @test x == 47
+        @test z == "z6"
+        @unpack_fields y = Property(47, 6.0, "z6")
+        @test y == 6.0
     end
 
     @testset "Type inference" begin
-        function f(z)
-            @unpack y, x = z
-            return x, y
+        function f(a)
+            @unpack y, z, x = a
+            return x, y, z
+        end
+        @test @inferred(f((; y=1.0, z="a", x=42))) == (42, 1.0, "a")
+        @test @inferred(f(Struct(42, 1.0, "a"))) == (42, 1.0, "a")
+        @test @inferred(f(Property(42, 1.0, "a"))) == (42, 1.0, "z")
+
+        function g(a)
+            @unpack_fields y, z, x = a
+            return x, y, z
         end
-        @test @inferred(f((; y=1.0, z="z", x=42))) == (42, 1.0)
-        @test @inferred(f(Struct(42, 1.0, "z"))) == (42, 1.0)
-        @test @inferred(f(Property(42, 1.0))) == (42, 1.0)
+        @test @inferred(g((; y=1.0, z="a", x=42))) == (42, 1.0, "a")
+        @test @inferred(g(Struct(42, 1.0, "a"))) == (42, 1.0, "a")
+        @test @inferred(g(Property(42, 1.0, "a"))) == (42, 1.0, "a")
     end
 
     @testset "Errors" begin
@@ -107,5 +155,10 @@ end
         @test_macro_throws ArgumentError @unpack (; x=42, y=1.0)
         @test_macro_throws ArgumentError @unpack x, y, (; x=42, y=1.0)
         @test_macro_throws ArgumentError @unpack x, 1 = (; x=42, y=1.0)
+
+        @test_macro_throws ArgumentError @unpack_fields d
+        @test_macro_throws ArgumentError @unpack_fields (; x=42, y=1.0)
+        @test_macro_throws ArgumentError @unpack_fields x, y, (; x=42, y=1.0)
+        @test_macro_throws ArgumentError @unpack_fields x, 1 = (; x=42, y=1.0)
     end
 end
