From 0a820c80b50ae5aefb65f0eaab62fb8395890e39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Fri, 18 Nov 2022 14:05:25 +0100
Subject: [PATCH] Add support for renamed local variables.

Fixes #18.

Incidental changes:

- wrap some tests into testsets, so that we can use `@isdefined` in tests

- add more validation of input expressions, test for expressions that
  are invalid but were expanded and failed with an obscure error message
---
 README.md        | 14 ++++++++++++++
 src/UnPack.jl    | 28 +++++++++++++++++++++-------
 test/runtests.jl | 48 ++++++++++++++++++++++++++++++++++--------------
 3 files changed, 69 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index d959b38..bb4e9a8 100644
--- a/README.md
+++ b/README.md
@@ -49,6 +49,20 @@ d = Dict{String,Any}()
 d # -> Dict{String,Any}("a"=>5.0,"c"=>"Hi!")
 ```
 
+Using `=>` allows unpacking to local variables that are different from a key:
+```julia
+struct MyContainer{T}
+    a::T
+    b::T
+end
+
+function Base.:(==)(x::MyContainer, y::MyContainer)
+    @unpack a, b = x
+    @unpack a => ay, b => by = y
+    a == ay && b ≈ by
+end
+```
+
 ## Customization of `@unpack` and `@pack!`
 
 What happens during the (un-)packing of a particular datatype is
diff --git a/src/UnPack.jl b/src/UnPack.jl
index 40f1961..4768c1e 100644
--- a/src/UnPack.jl
+++ b/src/UnPack.jl
@@ -81,20 +81,35 @@ Example with type:
 ```julia
 struct A; a; b; c; end
 d = A(4,7.0,"Hi")
-@unpack a, c = d
-a == 4 #true
-c == "Hi" #true
+@unpack a, c => C = d
+a == 4 # true
+C == "Hi" # note renaming, true
 ```
 
 Note that its functionality can be extended by adding methods to the
 `UnPack.unpack` function.
 """
 macro unpack(args)
-    args.head!=:(=) && error("Expression needs to be of form `a, b = c`")
+    (args isa Expr && args.head == :(=)) ||
+        error("Expression needs to be of form `a, b => b_renamed = c`")
     items, suitecase = args.args
-    items = isa(items, Symbol) ? [items] : items.args
+    items = (items isa Expr && items.head == :tuple) ? items.args : [items]
+    function _is_rename_expr(item)
+        (item isa Expr && item.head == :call) || return false
+        a = item.args
+        a[1] == :(=>) && a[2] isa Symbol && a[3] isa Symbol
+    end
     suitecase_instance = gensym()
-    kd = [:( $key = $UnPack.unpack($suitecase_instance, Val{$(Expr(:quote, key))}()) ) for key in items]
+    kd = map(items) do item
+        key, var = if item isa Symbol
+            item, item
+        elseif _is_rename_expr(item)
+            item.args[2], item.args[3]
+        else
+            error("Unrecognized key $(item). Keys need to be of the form `key` or `key => var`.")
+        end
+        :( $var = $UnPack.unpack($suitecase_instance, Val{$(Expr(:quote, key))}()) )
+    end
     kdblock = Expr(:block, kd...)
     expr = quote
         local $suitecase_instance = $suitecase # handles if suitecase is not a variable but an expression
@@ -104,7 +119,6 @@ macro unpack(args)
     esc(expr)
 end
 
-
 """
 ```julia_skip
 @pack! dict_or_typeinstance = a, b, c, ...
diff --git a/test/runtests.jl b/test/runtests.jl
index 31cec87..f602931 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -6,22 +6,42 @@ using Test
     ###########################
     # Packing and unpacking @unpack, @pack!
     ##########################
-    # Example with dict:
-    d = Dict{Symbol,Any}(:a=>5.0,:b=>2,:c=>"Hi!")
-    @unpack a, c = d
-    @test a == 5.0 #true
-    @test c == "Hi!" #true
 
-    d = Dict("a"=>5.0,"b"=>2,"c"=>"Hi!")
-    @unpack a, c = d
-    @test a == 5.0 #true
-    @test c == "Hi!" #true
+    @testset "dict with symbols" begin
+        d = Dict{Symbol,Any}(:a=>5.0,:b=>2,:c=>"Hi!")
+        @unpack a, c = d
+        @test a == 5.0 #true
+        @test c == "Hi!" #true
+    end
 
-    # Example with named tuple
-    @eval d = (a=5.0, b=2, c="Hi!")
-    @unpack a, c = d
-    @test a == 5.0 #true
-    @test c == "Hi!" #true
+    @testset "dict with strings" begin
+        d = Dict("a"=>5.0,"b"=>2,"c"=>"Hi!")
+        @unpack a, c = d
+        @test a == 5.0 #true
+        @test c == "Hi!" #true
+    end
+
+    @testset "named tuple" begin
+        @eval d = (a=5.0, b=2, c="Hi!")
+        @unpack a, c = d
+        @test a == 5.0 #true
+        @test c == "Hi!" #true
+    end
+
+    @testset "named tuple, renaming"  begin
+        d = (a = 1, b = 2)
+        @unpack a => c, b = d
+        @test c == 1
+        @test b == 2
+        @test !@isdefined a
+    end
+
+    @testset "invalid patterns" begin
+        @test_throws ErrorException macroexpand(Main, :(@unpack a))
+        @test_throws ErrorException macroexpand(Main, :(@unpack "a fish" = 1))
+        @test_throws ErrorException macroexpand(Main, :(@unpack a => "a fish" = 1))
+        @test_throws ErrorException macroexpand(Main, :(@unpack 1 => b = 1))
+    end
 end
 
 # having struct-defs inside a testset seems to be problematic in some julia version