Merge pull request #7 from avik-pal/ap/generalized

Adds a partial function macro `@$`

Author: archermarx

Project.toml

@@ -1,9 +1,13 @@
 name = "PartialFunctions"
 uuid = "570af359-4316-4cb7-8c74-252c00c2016b"
 authors = ["Thomas Marks <[email protected]> and contributors"]
-version = "1.1.1"
+version = "1.2.0"
+
+[deps]
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
 [compat]
+MacroTools = "0.5"
 julia = "1.5"
 
 [extras]

README.md

@@ -125,3 +125,20 @@ isequal(1, 2, ...)
 julia> isequal $ (1, 2) <| ()   # equivalent to a() or isequal(1, 2)
 false
 ```
+
+## The `@$` Macro
+
+`@$` allows users to create general partial functions by replacing the currently unknown
+arguments with `_`. For example, we can implement matrix multiplication as:
+
+```julia
+julia> matmul(A, X, B) = A * X .+ B
+
+julia> A = randn(2, 2); B = rand(2, 2); X = randn(2, 2);
+
+julia> pf = @$ matmul(_, X, _)
+matmul(_, X, _)
+
+julia> pf(A, B) ≈ matmul(A, X, B)
+true
+```

docs/src/index.md

@@ -137,3 +137,20 @@ julia> isequal $ (1, 2) <| ()   # equivalent to a() or isequal(1, 2)
 false
 ```
 
+## The `@$` Macro
+
+`@$` allows users to create general partial functions by replacing the currently unknown
+arguments with `_`. For example, we can implement matrix multiplication as:
+
+```jldoctest
+julia> matmul(A, X, B) = A * X .+ B
+matmul (generic function with 1 method)
+
+julia> A = randn(2, 2); B = rand(2, 2); X = randn(2, 2);
+
+julia> pf = @$ matmul(_, X, _)
+matmul(_, X, _)
+
+julia> pf(A, B) ≈ matmul(A, X, B)
+true
+```

src/PartialFunctions.jl

@@ -1,18 +1,26 @@
 module PartialFunctions
 
-export $
+using MacroTools
+
+export $, @$
 export <|
 include("reversedfunctions.jl")
 
 name(x) = (string ∘ Symbol)(x)
 
-struct PartialFunction{F<:Function, T<:Tuple, N<:NamedTuple} <: Function
+struct PartialFunction{KL, UL, F<:Function, T<:Tuple, N<:NamedTuple} <: Function
+    expr_string::String
     func::F
     args::T
     kwargs::N
-    PartialFunction(f, args::T, kwargs::N) where {T<:Tuple, N<:NamedTuple} = let
-        new{typeof(f), T, N}(f, args, kwargs)
-    end
+end
+
+function PartialFunction(f, args::T, kwargs::N) where {T<:Tuple, N<:NamedTuple}
+    return PartialFunction{nothing, nothing, typeof(f), typeof(args), typeof(kwargs)}("", f, args, kwargs)
+end
+
+function PartialFunction(known_args_locations, unknown_args_locations, expr_string::String, f, args::T, kwargs::N) where {T<:Tuple, N<:NamedTuple}
+    return PartialFunction{known_args_locations, unknown_args_locations, typeof(f), T, N}(expr_string, f, args, kwargs)
 end
 
 function PartialFunction(f::PartialFunction, newargs::Tuple, newkwargs::NamedTuple)
@@ -21,13 +29,13 @@ function PartialFunction(f::PartialFunction, newargs::Tuple, newkwargs::NamedTup
     PartialFunction(f.func, allargs, allkwargs)
 end
 
-(p::PartialFunction)(newargs...; newkwargs...) = p.func(p.args..., newargs...; p.kwargs..., newkwargs...)
+(p::PartialFunction{nothing})(newargs...; newkwargs...) = p.func(p.args..., newargs...; p.kwargs..., newkwargs...)
 
 """
     (\$)(f::Function, args...)
 Partially apply the given arguments to f. Typically used as infix `f \$ args`
 
-The returned function is of type [`PartialFunctions.PartialFunction{typeof(f), typeof(args)}`](@ref)
+The returned function is of type [`PartialFunctions.PartialFunction{nothing, nothing, typeof(f), typeof(args)}`](@ref)
 
 # Examples
 
@@ -52,6 +60,87 @@ end
 end
 
 ($)(f::DataType, args) = ($)(identity∘f, args)
+
+@generated function (pf::PartialFunction{kloc, uloc})(args...; kwargs...) where {kloc, uloc}
+    L = length(args)
+    if L == length(uloc)
+        # Execute the function
+        final_args = []
+        total_args = length(uloc) + length(kloc)
+        j, k = 1, 1
+        for i in 1:total_args
+            if i ∈ uloc
+                push!(final_args, :(args[$j]))
+                j += 1
+            else
+                push!(final_args, :(pf.args[$k]))
+                k += 1
+            end
+        end
+        return quote
+            pf.func($(final_args...); pf.kwargs..., kwargs...)
+        end
+    else
+        return :(pf $ (args..., (; kwargs...)))
+    end
+end
+
+
+"""
+    @\$ f(args...; kwargs...)
+
+Partially apply the given arguments to `f`. Unknown arguments are represented by `_`.
+
+!!! note
+    If no `_` is present, the function is executed immediately.
+
+# Examples
+
+```jldoctest
+julia> matmul(A, X, B; C = 1) = A * X .+ B .* C
+matmul (generic function with 1 method)
+
+julia> A = randn(2, 2); B = rand(2, 2); X = randn(2, 2);
+
+julia> pf = @\$ matmul(_, X, _; C = 2)
+matmul(_, X, _; C = 2)
+
+julia> pf(A, B) ≈ matmul(A, X, B; C = 2)
+true
+```
+"""
+macro ($)(expr::Expr)
+    if !@capture(expr, f_(args__; kwargs__) | f_(args__))
+        throw(ArgumentError("Only function calls are supported!"))
+    end
+
+    kwargs = kwargs === nothing ? [] : kwargs
+
+    underscore_args_pos = Tuple(findall(x -> x == :_, args))
+    if length(args) == 0 || length(underscore_args_pos) == 0
+        return :($(esc(expr)))
+    end
+
+    kwargs_keys = Symbol[]
+    kwargs_values = Any[]
+    for kwarg in kwargs
+        @assert kwarg.head == :kw "Malformed keyword argument!"
+        push!(kwargs_keys, kwarg.args[1])
+        push!(kwargs_values, kwarg.args[2])
+    end
+    kwargs_keys = Tuple(kwargs_keys)
+
+    non_underscore_args_pos = Tuple(setdiff(1:length(args), underscore_args_pos))
+    non_underscore_args = map(Base.Fix1(getindex, args), non_underscore_args_pos)
+    stored_args = NamedTuple{Symbol.(non_underscore_args_pos)}(non_underscore_args)
+
+    return :(PartialFunction($(esc(non_underscore_args_pos)),
+                             $(esc(underscore_args_pos)),
+                             $(esc(string(expr))), $(esc(f)),
+                             $(esc(tuple))($(esc.(non_underscore_args)...)),
+                             $(NamedTuple{kwargs_keys})(tuple($(esc.(kwargs_values)...)))))
+end
+
 """
     <|(f, args)
 
@@ -90,7 +179,9 @@ function Base.Symbol(pf::PartialFunction)
     Base.Symbol("$(func_name)$(argstring)")
 end
 
-Base.show(io::IO, pf::PartialFunction) = print(io, name(pf))
-Base.show(io::IO, ::MIME"text/plain", pf::PartialFunction) = show(io, pf)
+Base.show(io::IO, pf::PartialFunction{nothing}) = print(io, name(pf))
+Base.show(io::IO, ::MIME"text/plain", pf::PartialFunction{nothing}) = show(io, pf)
+Base.show(io::IO, pf::PartialFunction) = print(io, pf.expr_string)
+Base.show(io::IO, ::MIME"text/plain", pf::PartialFunction) = print(io, pf.expr_string)
 
 end

test/runtests.jl

@@ -4,49 +4,87 @@ using Test
 a(x) = x^2
 greet(greeting, name, punctuation) = "$(greeting), $(name)$(punctuation)"
 
-@testset "Partial functions" begin
-    @test map((+)$2, [1,2,3]) == [3, 4, 5]
-    @test repr(map $ a) == "map(a, ...)"
-    @test (map $ a)([1, 2, 3]) == [1, 4, 9]
-    
-    @test greet("Hello", "Bob", "!") == "Hello, Bob!"
-    sayhello = greet $ "Hello"
-    @test repr(sayhello) == "greet(\"Hello\", ...)"
-    @test repr("text/plain", sayhello) == repr(sayhello)
-
-    @test sayhello("Bob", "!") == "Hello, Bob!"
-    hi_bob = greet $ "Hi" $ "Bob" $ "!"
-    @test hi_bob isa PartialFunctions.PartialFunction{typeof(greet), Tuple{String, String, String}, NamedTuple{(), Tuple{}}}
-    @test hi_bob <| () == "Hi, Bob!"
-    @test sayhello <| ("Jimmy", "?")... == "Hello, Jimmy?"
-
-    @test greet $ ("Hi", "Bob") <| "!" == "Hi, Bob!"
-end
+@testset "PartialFunctions.jl" begin
+    @testset "Partial functions" begin
+        @test map((+)$2, [1,2,3]) == [3, 4, 5]
+        @test repr(map $ a) == "map(a, ...)"
+        @test (map $ a)([1, 2, 3]) == [1, 4, 9]
+        
+        @test greet("Hello", "Bob", "!") == "Hello, Bob!"
+        sayhello = greet $ "Hello"
+        @test repr(sayhello) == "greet(\"Hello\", ...)"
+        @test repr("text/plain", sayhello) == repr(sayhello)
 
-@testset "Reversed functions" begin
-    revmap = flip(map)
-    @test flip(revmap) == map
-    @test revmap([1,2,3], sin) == map(sin, [1,2,3])
-
-    func(x, y) = x - y
-    func(x, y, z) = x - y - z 
-    @test func(1, 2) == -1
-    @test func(1, 3, 6) == -8
-    flipped = flip(func)
-    @test flipped(2, 1) == -1
-    @test flipped(3, 2, 1) == -4
-end
+        @test sayhello("Bob", "!") == "Hello, Bob!"
+        hi_bob = greet $ "Hi" $ "Bob" $ "!"
+        @test hi_bob isa PartialFunctions.PartialFunction{nothing, nothing, typeof(greet), Tuple{String, String, String}, NamedTuple{(), Tuple{}}}
+        @test hi_bob <| () == "Hi, Bob!"
+        @test sayhello <| ("Jimmy", "?")... == "Hello, Jimmy?"
+
+        @test greet $ ("Hi", "Bob") <| "!" == "Hi, Bob!"
+    end
+
+    @testset "Reversed functions" begin
+        revmap = flip(map)
+        @test flip(revmap) == map
+        @test revmap([1,2,3], sin) == map(sin, [1,2,3])
+
+        func(x, y) = x - y
+        func(x, y, z) = x - y - z 
+        @test func(1, 2) == -1
+        @test func(1, 3, 6) == -8
+        flipped = flip(func)
+        @test flipped(2, 1) == -1
+        @test flipped(3, 2, 1) == -4
+    end
+
+    @testset "Keyword Arguments" begin
+        a = [[1,2,3], [1,2]]
+        sort_by_length = sort $ (; by = length)
+        @test sort(a, by = length) == sort_by_length(a)
+
+        sort_a_by_length = sort $ (a, (;by = length))
+        @test sort(a, by = length) == sort_a_by_length()
 
-@testset "Keyword Arguments" begin
-    a = [[1,2,3], [1,2]]
-    sort_by_length = sort $ (; by = length)
-    @test sort(a, by = length) == sort_by_length(a)
+        sort_a_by_length_2 = sort $ ((a,), (;by = length))
+        @test sort_a_by_length == sort_a_by_length_2
 
-    sort_a_by_length = sort $ (a, (;by = length))
-    @test sort(a, by = length) == sort_a_by_length()
+        @test repr(sort_a_by_length) == "sort([[1, 2, 3], [1, 2]], ...; by = length, ...)"
+    end
 
-    sort_a_by_length_2 = sort $ ((a,), (;by = length))
-    @test sort_a_by_length == sort_a_by_length_2
+    @testset "Generalized Partial Functions" begin
+        @test map(@$(+(2, _)), [1,2,3]) == [3, 4, 5]
+        @test map(@$(+(_, 2)), [1,2,3]) == [3, 4, 5]
+        @test repr(@$(map(a, _))) == "map(a, _)"
+        @test (@$(map(a, _)))([1, 2, 3]) == [1, 4, 9]
+        
+        @test greet("Hello", "Bob", "!") == "Hello, Bob!"
+        sayhello = @$ greet("Hello", _, _)
+        @test repr(sayhello) == "greet(\"Hello\", _, _)"
+        @test repr("text/plain", sayhello) == repr(sayhello)
 
-    @test repr(sort_a_by_length) == "sort([[1, 2, 3], [1, 2]], ...; by = length, ...)"
-end
+        @test sayhello("Bob", "!") == "Hello, Bob!"
+
+        sayhellobob = sayhello("Bob")
+        @test repr(sayhellobob) == "greet(\"Hello\", \"Bob\", ...)"
+        @test sayhellobob("!") == "Hello, Bob!"
+
+        hi_bob = @$(@$(greet("Hi", _, _))("Bob", _))
+        @test @$(hi_bob("!")) == "Hi, Bob!"
+        @test hi_bob isa PartialFunctions.PartialFunction
+        @test sayhello <| ("Jimmy", "?")... == "Hello, Jimmy?"
+
+        @test hi_bob <| "!" == "Hi, Bob!"
+
+        @testset "Keyword Arguments" begin
+            a = [[1,2,3], [1,2]]
+            sort_by_length = @$(sort(_; by = length))
+            @test sort(a, by = length) == sort_by_length(a)
+
+            sorted_a_by_length = @$(sort(a; by = length))
+            @test sort(a, by = length) == sorted_a_by_length
+
+            @test repr(sort_by_length) == "sort(_; by = length)"
+        end
+    end
+end