Recurse _ to bypass square brackets, and usually-infix operators

src/Underscores.jl

@@ -62,6 +62,38 @@ end
 replace_(ex)  = add_closures(ex, "_", r"^_([0-9]*|[₀-₉]*)$")
 replace__(ex) = add_closures(ex, "__", r"^__([0-9]*|[₀-₉]*)$")
 
+const _square_bracket_ops = [:comprehension, :typed_comprehension, :generator,
+                             :ref, :vcat, :typed_vcat, :hcat, :typed_hcat, :row]
+
+_isoperator(x) = x isa Symbol && Base.isoperator(x)
+
+function lower_inner(ex)
+    if ex isa Expr
+        if ex.head == :(=) ||
+            (ex.head == :call && length(ex.args) > 1 && _isoperator(ex.args[1]))
+            # Infix operators do not count as outermost function call
+            return Expr(ex.head, ex.args[1],
+                map(lower_inner, ex.args[2:end])...)
+        elseif ex.head in _square_bracket_ops
+            # Indexing & other square brackets not counted as outermost function
+            return Expr(ex.head, map(lower_inner, ex.args)...)
+        elseif ex.head == :. && length(ex.args) == 2 && ex.args[2] isa QuoteNode
+            # Getproperty also doesn't count
+            return Expr(ex.head, map(lower_inner, ex.args)...)
+        elseif ex.head == :. && length(ex.args) == 2 &&
+            ex.args[2] isa Expr && ex.args[2].head == :tuple
+            # Broadcast calls treated as normal calls for underscore lowering
+            return Expr(ex.head, replace_(ex.args[1]),
+                                  Expr(:tuple, map(replace_, ex.args[2].args)...))
+        else
+            # For other syntax, replace _ in args individually
+            return Expr(ex.head, map(replace_, ex.args)...)
+        end
+    else
+        return ex
+    end
+end
+
 # In principle this can be extended locally by a package for use within the
 # package and for prototyping purposes. However note that this will interact
 # badly with precompilation. (If it makes sense we could fix this per-package
@@ -77,17 +109,11 @@ function lower_underscores(ex)
             # Special case for pipelining and composition operators
             return Expr(ex.head, ex.args[1],
                         map(lower_underscores, ex.args[2:end])...)
-        elseif ex.head == :.       && length(ex.args) == 2 &&
-               ex.args[2] isa Expr && ex.args[2].head == :tuple
-            # Broadcast calls treated as normal calls for underscore lowering
-            return replace__(Expr(ex.head, replace_(ex.args[1]),
-                                  Expr(:tuple, map(replace_, ex.args[2].args)...)))
         elseif ex.head == :do
             error("@_ expansion for `do` syntax is reserved")
         else
-            # For other syntax, replace _ in args individually and __ over the
-            # entire expression.
-            return replace__(Expr(ex.head, map(replace_, ex.args)...))
+            # For other syntax, replace __ over the entire expression
+            return replace__(lower_inner(ex))
         end
     else
         return ex
@@ -102,7 +128,7 @@ and *pass them along* to `func`.
 
 The detailed rules are:
 1. Uses of the placeholder `_` expand to the single argument of an anonymous
-   function which is passed to the outermost expression.
+   function which is passed to the outermost ordinary function call.
 2. Numbered placeholders `_1,_2,...` (or `_₁,_₂,...`) may be used if you need
    more than one argument. Numbers indicate position in the argument list.
 3. The double underscore placeholder `__` (and numbered versions `__1,__2,...`)
@@ -162,6 +188,36 @@ julia> @_ table |>
  2
  3
 ```
+
+## Extraordinary functions
+
+The scope of `_` as described in rule 1 depends on "ordinary" function call.
+This excludes the following operations:
+
+* Square brackets: In `map(_^2,__)[3]`, it is `map` which receives an anonymous
+  function, as this happens before the indexing is lowered to `getindex(...,3)`.
+  Comprehensions (`collect`) and explicit matrix constructions (`hvcat`) are
+  treated similarly.
+
+* Broadcasting, and field access: In `f.(_,xs)` and `f(_,x).y` the function `f`
+  is the ordinary call, not the internal `broadcast` or `getproperty`.
+
+* Infix operators: While `sum(_^2,x) / length(x)` can be written in prefix form
+  `/(...,...)`, the convention of `@_` is not to view this as an ordinary call,
+  and hence to pass the anonymous function to `sum` instead. This also applies to
+  broadcasted operators, such as `map(_^2,x) ./ length(x)`.
+
+The scope of `__` is unaffected by these concerns.
+
+| Expression                       | Meaning                            |
+|:-------------------------------- |:---------------------------------- |
+| `@_ data \\|> map(_[2],__)[3]`   | `data \\|> (d->map(x->x[2],d)[3])` |
+| `@_ [sum(_*_, z) for z in a]`    | `[sum(x->x*x, z) for z in a]`      |
+| `@_ sum(1+_^2, data).re`         | `sum(x->1+x^2, data).re`           |
+| `@_ sum(_^2,a) / length(a)`      | `sum(x->x^2,a) / length(a)`        |
+| `@_ /(sum(_^2,a), length(a))`    | The same, infix form is canonical. |
+| `@_ data \\|> filter(_>3,__).^2` | `data \\|> d->(filter(>(3),d).^2)` |
+
 """
 macro _(ex)
     esc(lower_underscores(ex))

test/runtests.jl

@@ -14,6 +14,16 @@ using Test
     @test data[1:1] == @_ filter(startswith(_.x, "a"), data)
     @test data[2:3] == @_ filter(_.y >= 2, data)
 
+    # Use with indexing
+    @test data[1] == @_ filter(startswith(_.x, "a"), data)[end]
+    @test data[2:3] == @_ filter(_.y >= 2, data)[1:2]
+    @test data[2] == @_ data[findfirst(_.y >= 2, data)]
+
+    # Operators
+    @test -2 == @_ -findfirst(_.x == "b", data)
+    @test [14] == @_ [1 2 3] * map(_.y, data)
+    @test (1:3)./3 == @_ data |> map(_.y, __) ./ length(__)
+
     # Multiple args
     @test [0,0] == @_ map(_-_, [1,2])
 
@@ -31,6 +41,10 @@ using Test
 
     @test [0,0,0] == @_ data |> map(_1.y + _2, __, [-1,-2,-3])
 
+    @test 3 == @_ [[1], [1,2], [1,2,3]] |>
+                                map(_[_[end]], __) |>
+                                __[end]
+
     # Use with piping and lazy versions of map and filter
     Filter(f) = x->filter(f,x)
     Map(f) = x->map(f,x)
@@ -40,6 +54,28 @@ using Test
                     Map(_.y)
 
     @test [1] == @_(Map(_.y) ∘ Filter(startswith(_.x, "a")))(data)
+
+    # Getproperty
+    @test "a" == @_ data |> __[1].x
+    @test "b" == @_ (x="a", y=(z="b", t="c")) |> __.y.z
+    @test "c" == @_ (x="a", y=(z="b", t="c")) |> __.y[end]
+    @test -3 == @_ [1+2im,3,4+5im] |> sum(_^2, __).re
+
+    # Interpolation
+    @test ["1","a","2.0"] == @_ map("$_", [1,"a",2.0])
+
+    # Broadcasting
+    @_ [[1],[2],[3]] == data |> filter.(_ isa Int, collect.(__))
+
+    # Comprehensions
+    @test [[],[],[3]] == @_ [[1], [1,2], [1,2,3]] |>
+                                [filter(_>2, x) for x in __]
+    @test fill(:y,3) == @_ data |>
+                        Symbol[findfirst(_ isa Int, x) for x in __]
+
+    # Matrix construction
+    @test 4*ones(2,2) == @_ ones(4) |> [ sum(__)      sum(_^2, __)
+                                         sum(_^3, __) sum(_^4, __) ]
 end
 
 @testset "Underscores lowering" begin
@@ -114,6 +150,11 @@ end
     @test lower(:(f.(__)))  == cleanup!(:((__1,)->f.(__1)))
     @test lower(:((_).(x)))  == cleanup!(:(((_1,)->_1).(x)))
 
+    # Indexing
+    @test lower(:(f(_)[2])) == cleanup!(:(f((_1,)->_1)[2]))
+    @test lower(:(f(g(_[3]),h)[4])) == cleanup!(:(f((_1,)->g(_1[3]),h)[4]))
+    @test lower(:(f(__)[5])) == cleanup!(:((__1,)->f(__1)[5]))
+
     # Random sample of other syntax
     @test lower(:([_]))  == cleanup!(:([(_1,)->_1]))
     @test lower(:((f(_), g(_))))  == cleanup!(:(((_1,)->f(_1)), ((_1,)->g(_1))))