add KeyPath function

add to docs

fmap_with_keypath

simplify names

add fmapstructure_with_path

update fmapstructure to handle multiple arguments

fixes

Update src/Functors.jl

Co-authored-by: Kyle Daruwalla <[email protected]>

Update src/Functors.jl

Co-authored-by: Kyle Daruwalla <[email protected]>

Update src/Functors.jl

Co-authored-by: Kyle Daruwalla <[email protected]>

Author: CarloLucibello

docs/src/api.md

@@ -1,5 +1,6 @@
 ```@docs
 Functors.fmap
+Functors.fmap_with_path
 Functors.@functor
 Functors.@leaf
 ```
@@ -24,6 +25,11 @@ Functors.IterateWalk
 
 ```@docs
 Functors.fmapstructure
+Functors.fmapstructure_with_path
 Functors.fcollect
 Functors.fleaves
 ```
+
+```@docs
+Functors.KeyPath
+```

src/Functors.jl

@@ -1,8 +1,10 @@
 module Functors
 
-export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect, execute, fleaves
+export @functor, @flexiblefunctor, fmap, fmapstructure, fcollect, execute, fleaves, 
+       KeyPath, fmap_with_path, fmapstructure_with_path
 
 include("functor.jl")
+include("keypath.jl")
 include("walks.jl")
 include("maps.jl")
 include("base.jl")
@@ -104,7 +106,7 @@ Equivalent to `functor(x)[1]`.
 children
 
 """
-    fmap(f, x, ys...; exclude = Functors.isleaf, walk = Functors.DefaultWalk()[, prune])
+    fmap(f, x, ys...; exclude = Functors.isleaf, walk = Functors.DefaultWalk(), [prune])
 
 A structure and type preserving `map`.
 
@@ -113,6 +115,8 @@ by applying `f`, and otherwise traverses `x` recursively using [`functor`](@ref)
 Optionally, it may also be associated with objects `ys` with the same tree structure.
 In that case, `f` is applied to the corresponding leaf nodes in `x` and `ys`.
 
+See also [`fmap_with_path`](@ref) and [`fmapstructure`](@ref).
+
 # Examples
 ```jldoctest
 julia> fmap(string, (x=1, y=(2, 3)))
@@ -222,14 +226,16 @@ fmap
 
 
 """
-    fmapstructure(f, x; exclude = isleaf)
+    fmapstructure(f, x, ys...; exclude = isleaf, [prune])
 
 Like [`fmap`](@ref), but doesn't preserve the type of custom structs.
 Instead, it returns a `NamedTuple` (or a `Tuple`, or an array),
 or a nested set of these.
 
 Useful for when the output must not contain custom structs.
 
+See also [`fmap`](@ref) and [`fmapstructure_with_path`](@ref).
+
 # Examples
 ```jldoctest
 julia> struct Foo; x; y; end
@@ -304,6 +310,47 @@ julia> fcollect(m, exclude = v -> Functors.isleaf(v))
 fcollect
 
 
+""""
+    fmap_with_path(f, x, ys...; exclude = isleaf, walk = DefaultWalkWithPath(), [prune])
+
+Like [`fmap`](@ref), but also passes a [`KeyPath`](@ref Functors.KeyPath) to `f` for each node in the
+recursion. The `KeyPath` is a tuple of the indices used to reach the current
+node from the root of the recursion. The `KeyPath` is constructed by the
+`walk` function, and can be used to reconstruct the path to the current node
+from the root of the recursion.
+
+`f` has to accept two arguments: the associated `KeyPath` and the value of the current node.
+
+`exclude` also receives the `KeyPath` as its first argument and a node as its second.
+It should return `true` if the recursion should not continue on its children and `f` applied to it.
+
+`prune` is used to control the behaviour when the same node appears twice, see [`fmap`](@ref)
+for more information.
+
+# Examples
+
+```jldoctest
+julia> x = ([1, 2, 3], 4, (a=5, b=Dict("A"=>6, "B"=>7), c=Dict("C"=>8, "D"=>9)));
+
+julia> fexclude(kp, x) = kp == KeyPath(3, :c) || Functors.isleaf(x);
+
+julia> fmap_with_path((kp, x) -> x isa Dict ? nothing : x.^2, x; exclude = fexclude)
+([1, 4, 9], 16, (a = 25, b = Dict("B" => 49, "A" => 36), c = nothing))
+```
+"""
+fmap_with_path
+
+"""
+    fmapstructure_with_path(f, x, ys...; [exclude, prune])
+
+Like [`fmap_with_path`](@ref), but doesn't preserve the type of custom structs.
+Instead, it returns a named tuple, a tuple, an array, a dict, 
+or a nested set of these.
+
+See also [`fmapstructure`](@ref).
+"""
+fmapstructure_with_path
+
 
 """
     fleaves(x; exclude = isleaf)
@@ -332,7 +379,9 @@ julia> fleaves(m)
 2-element Vector{Any}:
  [1, 2, 3]
  TypeWithNoChildren(:a, :b)
+```
 """
 fleaves
 
 end # module
+

src/keypath.jl

@@ -0,0 +1,47 @@
+KeyT = Union{Symbol, AbstractString, Integer}
+
+"""
+    KeyPath(keys...)
+
+A type for representing a path of keys to a value in a nested structure.
+Can be constructed with a sequence of keys, or by concatenating other `KeyPath`s.
+Keys can be of type `Symbol`, `String`, or `Int`.
+
+# Examples
+
+```jldoctest
+julia> kp = KeyPath(:b, 3)
+KeyPath(:b, 3)
+
+julia> KeyPath(:a, kp, :c, 4)
+KeyPath(:a, :b, 3, :c, 4)
+```
+"""
+struct KeyPath{T<:Tuple}
+    keys::T    
+end
+
+@functor KeyPath
+isleaf(::KeyPath, @nospecialize(x)) = isleaf(x)
+
+function KeyPath(keys::Union{KeyT, KeyPath}...)
+    ks = (k isa KeyPath ? (k.keys...,) : (k,) for k in keys)
+    return KeyPath(((ks...)...,))
+end
+
+Base.getindex(kp::KeyPath, i::Int) = kp.keys[i]
+Base.length(kp::KeyPath) = length(kp.keys)
+Base.iterate(kp::KeyPath, state=1) = iterate(kp.keys, state)
+Base.:(==)(kp1::KeyPath, kp2::KeyPath) = kp1.keys == kp2.keys
+
+function Base.show(io::IO, kp::KeyPath)
+    compat = get(io, :compact, false)
+    if compat
+        print(io, keypathstr(kp))          
+    else
+        print(io, "KeyPath$(kp.keys)")
+    end
+end
+
+keypathstr(kp::KeyPath) = join(kp.keys, ".")
+

src/maps.jl

@@ -11,7 +11,22 @@ function fmap(f, x, ys...; exclude = isleaf,
   execute(_walk, x, ys...)
 end
 
-fmapstructure(f, x; kwargs...) = fmap(f, x; walk = StructuralWalk(), kwargs...)
+function fmap_with_path(f, x, ys...; 
+                exclude = isleaf,
+                walk = DefaultWalkWithPath(),
+                cache = IdDict(),
+                prune = NoKeyword())
+  
+  _walk = ExcludeWalkWithKeyPath(walk, f, exclude)
+  if !isnothing(cache)
+    _walk = CachedWalkWithPath(_walk, prune, cache)
+  end
+  return execute(_walk, KeyPath(), x, ys...)
+end
+
+fmapstructure(f, x, ys...; kwargs...) = fmap(f, x, ys...; walk = StructuralWalk(), kwargs...)
+
+fmapstructure_with_path(f, x, ys...; kwargs...) = fmap_with_path(f, x, ys...; walk = StructuralWalkWithPath(), kwargs...)
 
 fcollect(x; exclude = v -> false) =
   execute(ExcludeWalk(CollectWalk(), _ -> nothing, exclude), x)

src/walks.jl

@@ -1,9 +1,24 @@
-_map(f, x...) = map(f, x...)
+function _map(f, x, ys...)
+  check_lenghts(x, ys...) || error("all arguments must have at least the same length of the firs one")
+  map(f, x, ys...)
+end
+
+function check_lenghts(x, ys...)
+  n = length(x)
+  return all(y -> length(y) >= n, ys)
+end
+
 _map(f, x::Dict, ys...) = Dict(k => f(v, (y[k] for y in ys)...) for (k, v) in x)
 
 _values(x) = x
 _values(x::Dict) = values(x)
 
+_keys(x::Dict) = Dict(k => k for k in keys(x))
+_keys(x::Tuple) = (keys(x)...,)
+_keys(x::AbstractArray) = collect(keys(x))
+_keys(x::NamedTuple{Ks}) where Ks = NamedTuple{Ks}(Ks)
+
+
 """
     AbstractWalk
 
@@ -76,6 +91,16 @@ function (::DefaultWalk)(recurse, x, ys...)
   re(_map(recurse, func, yfuncs...))
 end
 
+struct DefaultWalkWithPath <: AbstractWalk end
+
+function (::DefaultWalkWithPath)(recurse, kp::KeyPath, x, ys...)
+  x_children, re = functor(x)
+  kps = _map(c -> KeyPath(kp, c), _keys(x_children)) # use _keys and _map to preserve x_children type
+  ys_children = map(children, ys)
+  re(_map(recurse, kps, x_children, ys_children...))
+end
+
+
 """
     StructuralWalk()
 
@@ -86,7 +111,20 @@ See [`fmapstructure`](@ref) for more information.
 """
 struct StructuralWalk <: AbstractWalk end
 
-(::StructuralWalk)(recurse, x) = _map(recurse, children(x))
+function (::StructuralWalk)(recurse, x, ys...)
+  x_children = children(x)
+  ys_children = map(children, ys)
+  return _map(recurse, x_children, ys_children...)
+end
+
+struct StructuralWalkWithPath <: AbstractWalk end
+
+function (::StructuralWalkWithPath)(recurse, kp::KeyPath, x, ys...)
+  x_children = children(x)
+  kps = _map(c -> KeyPath(kp, c), _keys(x_children)) # use _keys and _map to preserve x_children type
+  ys_children = map(children, ys)
+  return _map(recurse, kps, x_children, ys_children...)
+end
 
 """
     ExcludeWalk(walk, fn, exclude)
@@ -106,6 +144,16 @@ end
 (walk::ExcludeWalk)(recurse, x, ys...) =
   walk.exclude(x) ? walk.fn(x, ys...) : walk.walk(recurse, x, ys...)
 
+struct ExcludeWalkWithKeyPath{T, F, G} <: AbstractWalk
+  walk::T
+  fn::F
+  exclude::G
+end
+
+(walk::ExcludeWalkWithKeyPath)(recurse, kp::KeyPath, x, ys...) =
+  walk.exclude(kp, x) ? walk.fn(kp, x, ys...) : walk.walk(recurse, kp, x, ys...)
+  
+
 struct NoKeyword end
 
 usecache(::Union{AbstractDict, AbstractSet}, x) =
@@ -152,6 +200,28 @@ function (walk::CachedWalk)(recurse, x, ys...)
   end
 end
 
+struct CachedWalkWithPath{T, S} <: AbstractWalk
+  walk::T
+  prune::S
+  cache::IdDict{Any, Any}
+end
+
+CachedWalkWithPath(walk; prune = NoKeyword(), cache = IdDict()) =
+  CachedWalkWithPath(walk, prune, cache)
+
+function (walk::CachedWalkWithPath)(recurse, kp::KeyPath, x, ys...)
+  should_cache = usecache(walk.cache, x)
+  if should_cache && haskey(walk.cache, x)
+    return walk.prune isa NoKeyword ? walk.cache[x] : walk.prune
+  else
+    ret = walk.walk(recurse, kp, x, ys...)
+    if should_cache
+      walk.cache[x] = ret
+    end
+    return ret
+  end
+end
+
 """
     CollectWalk()