Merge pull request #10 from darsnack/darsnack/predicate-fmap

Add predicate keyword to fmap

Author: CarloLucibello

README.md

@@ -71,3 +71,22 @@ Any field not in the list will not be returned by `functor` and passed through a
 It is also possible to implement `functor` by hand when greater flexibility is required. See [here](https://github.com/FluxML/Functors.jl/issues/3) for an example.
 
 For a discussion regarding the need for a `cache` in the implementation of `fmap`, see [here](https://github.com/FluxML/Functors.jl/issues/2).
+
+Use `exclude` for more fine-grained control over whether `fmap` descends into a particular value (the default is `exclude = Functors.isleaf`):
+```julia
+julia> using CUDA
+
+julia> x = ['a', 'b', 'c'];
+
+julia> fmap(cu, x)
+3-element Array{Char,1}:
+ 'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
+ 'b': ASCII/Unicode U+0062 (category Ll: Letter, lowercase)
+ 'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
+
+julia> fmap(cu, x; exclude = x -> CUDA.isbitstype(eltype(x)))
+3-element CuArray{Char,1}:
+ 'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
+ 'b': ASCII/Unicode U+0062 (category Ll: Letter, lowercase)
+ 'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
+```

src/functor.jl

@@ -51,9 +51,12 @@ end
 
 # See https://github.com/FluxML/Functors.jl/issues/2 for a discussion regarding the need for
 # cache.
-function fmap(f, x; cache = IdDict())
+function fmap(f, x; exclude = isleaf, cache = IdDict())
   haskey(cache, x) && return cache[x]
-  cache[x] = isleaf(x) ? f(x) : fmap1(x -> fmap(f, x, cache = cache), x)
+  y = exclude(x) ? f(x) : fmap1(x -> fmap(f, x, cache = cache, exclude = exclude), x)
+  cache[x] = y
+
+  return y
 end
 
 """

test/basics.jl

@@ -32,7 +32,20 @@ end
   @test model′.x.y isa Vector{Float64}
 end
 
-@testset "Property list" begin  
+@testset "Exclude" begin
+  f(x::AbstractArray) = x
+  f(x::Char) = 'z'
+
+  x = ['a', 'b', 'c']
+  @test fmap(f, x)  == ['z', 'z', 'z']
+  @test fmap(f, x; exclude = x -> x isa AbstractArray) == x
+
+  x = (['a', 'b', 'c'], ['d', 'e', 'f'])
+  @test fmap(f, x)  == (['z', 'z', 'z'], ['z', 'z', 'z'])
+  @test fmap(f, x; exclude = x -> x isa AbstractArray) == x
+end
+
+@testset "Property list" begin
   model = Baz(1, 2, 3)
   model′ = fmap(x -> 2x, model)