fix mapexpr and further specialize on function call names for :call e…

…xprs

src/frompackage/code_parsing.jl

@@ -11,7 +11,10 @@ function map_and_clean_expr(ex::Expr, mapexpr)
  @nospecialize
  new_args = map(mapexpr, ex.args)
  filter!(x -> !isa(x, RemoveThisExpr), new_args)
- return isempty(new_args) ? RemoveThisExpr() : Expr(ex.head, new_args...)
+ # If we still have args, or if the head is a block, we return the modified haed. We have this special case for a block for tests mostly
+ !isempty(new_args) | (ex.head === :block) && return Expr(ex.head, new_args...)
+ # If we get here we just return RemoveThisExpr
+ return RemoveThisExpr()
 end
 
 ## custom_walk! ##
@@ -66,12 +69,18 @@ function custom_walk!(::AbstractEvalController, ex::Expr, ::Val{:struct})
  return ex
 end
 
+function custom_walk!(p::AbstractEvalController, ex::Expr, v::Val{:call})
+ @nospecialize
+ func_name = first(ex.args) |> Symbol
+ return custom_walk!(p, ex, v, Val{func_name}())
+end
+function custom_walk!(p::AbstractEvalController, ex::Expr, ::Val{:call}, ::Val)
+ @nospecialize
+ return map_and_clean_expr(ex, p.custom_walk)
+end
 # This handles include calls, by adding p.custom_walk as the modexpr
-function custom_walk!(p::AbstractEvalController, ex::Expr, ::Val{:call})
+function custom_walk!(p::AbstractEvalController, ex::Expr, ::Val{:call}, ::Val{:include})
  @nospecialize
- f = p.custom_walk
- # We just process this expression if it's not an `include` call
- first(ex.args) === :include || return Expr(:call, map(f, ex.args)...)
  new_ex = :($process_include_expr!($p))
  append!(new_ex.args, ex.args[2:end])
  return new_ex

test/TestPackage/out_notebook.jl

@@ -66,13 +66,13 @@ This tests that using custom `mapexpr` function within include statements in the
 """
 
 # ╔═╡ c3691cc1-9cfb-460e-b94a-15c9357b0892
-TestPackage.ModExpr.should_be_100 == 100 || error("The custom application of mapexpr in `include` did not work")
+TestPackage.MapExpr.should_be_100 == 100 || error("The custom application of mapexpr in `include` did not work")
 
 # ╔═╡ 83b8c9bb-c831-48cd-8d5d-0c9b8691d35c
-TestPackage.ModExpr.should_be_1 == 1 || error("The custom application of mapexpr in `include` did not work.\nVariable `should_be_1` is not 1")
+TestPackage.MapExpr.should_be_1 == 1 || error("The custom application of mapexpr in `include` did not work.\nVariable `should_be_1` is not 1")
 
 # ╔═╡ a298f325-85ef-4b48-a15c-67080515dd8d
-!isdefined(TestPackage.ModExpr, :var_to_delete) || error("The custom application of mapexpr in `include` did not work")
+!isdefined(TestPackage.MapExpr, :var_to_delete) || error("The custom application of mapexpr in `include` did not work")
 
 # ╔═╡ 79a035c0-9138-4b47-87ab-1ce5e4b9a4d4
 md"""

test/TestPackage/src/TestPackage.jl

@@ -43,11 +43,12 @@ module PrettyPrint
  struct SomeType end
 end
 
-module ModExpr
- import PlutoDevMacros.FromPackage: RemoveThisExpr
+module MapExpr
+ import PlutoDevMacros.FromPackage: RemoveThisExpr, map_and_clean_expr
  import PlutoDevMacros.FromPackage.MacroTools: postwalk
  function mapexpr(ex)
- Meta.isexpr(ex, :(=)) || return ex
+ ex isa Expr || return ex
+ Meta.isexpr(ex, :(=)) || return map_and_clean_expr(ex, identity)
  args = deepcopy(ex.args)
  varname = first(args)
  varname === :var_to_delete && return RemoveThisExpr()