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

src/Underscores.jl

@@ -68,7 +68,12 @@ replace__(ex) = add_closures(ex, "__", r"^__([0-9]*|[₀-₉]*)$")
 # by storing a per-module _pipeline_ops in the module using @_.)
 const _pipeline_ops = [:|>, :<|, :∘]
 
-function lower_underscores(ex)
+const _square_bracket_ops = [:comprehension, :typed_comprehension, :generator,
+                             :vcat, :typed_vcat, :hcat, :typed_hcat, :row]
+
+_isoperator(x) = x isa Symbol && Base.isoperator(x)
+
+function lower_underscores(ex, replace__=replace__)
     if ex isa Expr
         if isquoted(ex)
             return ex
@@ -77,11 +82,23 @@ 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 == :(=) ||
+            (ex.head == :call && length(ex.args) > 1 && _isoperator(ex.args[1]))
+            # Other operators do not count as outermost function call
+            return replace__(Expr(ex.head, ex.args[1],
+                map(x -> lower_underscores(x, identity), 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 == :ref && ex.args[1] isa Expr
+            # Indexing is not counted as outermost function
+            return replace__(Expr(ex.head,
+                Expr(ex.args[1].head, map(replace_, ex.args[1].args)...), ex.args[2]))
+        elseif ex.head in _square_bracket_ops
+            return replace__(Expr(ex.head,
+                                  map(x -> lower_underscores(x, identity), ex.args)...))
         elseif ex.head == :do
             error("@_ expansion for `do` syntax is reserved")
         else
@@ -102,7 +119,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 +179,33 @@ 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)`.
+  Similarly, generator expressions such as `[sum(_^2, x) for x in __]` are not
+  simply calls to `collect`, and constructions like `Int[sum(_^2, __) sum(_^3, __)]`
+  are not ordinary calls to `(hv)cat`.
+
+* 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(_^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,15 @@ 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]
+
+    # 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 +40,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 +53,19 @@ using Test
                     Map(_.y)
 
     @test [1] == @_(Map(_.y) ∘ Filter(startswith(_.x, "a")))(data)
+
+    # Interpolation
+    @test ["1","a","2.0"] == @_ map("$_", [1,"a",2.0])
+
+    # 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 +140,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))))