Add nitems, firstitem(s), and lastitem(s) (#37)

See also: JuliaLang/julia#35947

Author: tkf

src/DataTools.jl

@@ -1,18 +1,47 @@
 module DataTools
 
-export averaging, inc1, meanvar, modifying, oncol, rightif
+export averaging,
+    firstitem,
+    firstitems,
+    inc1,
+    lastitem,
+    lastitems,
+    meanvar,
+    modifying,
+    nitems,
+    oncol,
+    rightif
 
+using Base: HasLength, HasShape, IteratorSize
 using InitialValues: InitialValues
 using Setfield: @lens, Lens, PropertyLens, modify, set
 using StaticNumbers: static
 using Statistics: Statistics, mean, std, var
 using Tables: Tables
-using Transducers: Map, Transducers, combine, complete, next, reducingfunction, start
+using Transducers:
+    Composition,
+    Count,
+    IdentityTransducer,
+    Map,
+    MapSplat,
+    Scan,
+    Take,
+    TakeLast,
+    Transducers,
+    combine,
+    complete,
+    extract_transducer,
+    next,
+    opcompose,
+    reducingfunction,
+    right,
+    start
 
 include("utils.jl")
 include("oncol.jl")
 include("modifying.jl")
 include("reductions.jl")
+include("reducers.jl")
 
 # Use README as the docstring of the module:
 @doc let path = joinpath(dirname(@__DIR__), "README.md")

src/reducers.jl

@@ -0,0 +1,108 @@
+"""
+    nitems(xs) -> n::Integer
+
+Count number of items in `xs`.  Consume `xs` if necessary.
+
+# Examples
+```jldoctest
+julia> using DataTools, Transducers
+
+julia> nitems(1:10)
+10
+
+julia> 1:10 |> Filter(isodd) |> Map(inv) |> nitems
+5
+```
+"""
+nitems
+nitems(xs) =
+    if IteratorSize(xs) isa Union{HasLength, HasShape}
+        length(xs)
+    else
+        xf, coll = extract_transducer(xs)
+        _nitems(_pop_innermost_maplikes(xf), coll)
+    end
+
+_pop_innermost_maplikes(xf) = xf
+_pop_innermost_maplikes(::Union{Map,MapSplat,Scan}) = IdentityTransducer()
+function _pop_innermost_maplikes(xf::Composition)
+    inner = _pop_innermost_maplikes(xf.inner)
+    if inner isa IdentityTransducer
+        return _pop_innermost_maplikes(xf.outer)
+    else
+        opcompose(xf.outer, inner)
+    end
+end
+
+_nitems(::IdentityTransducer, xs) = _nitems(xs)
+_nitems(xf, xs) = xs |> xf |> _nitems
+
+_nitems(xs) =
+    if IteratorSize(xs) isa Union{HasLength, HasShape}
+        length(xs)
+    else
+        foldl(inc1, IdentityTransducer(), xs)
+    end
+# TODO: optimization for `Cat`.
+
+"""
+    firstitem(xs)
+
+Get the first item of `xs`.  Consume `xs` if necessary.
+
+# Examples
+```jldoctest
+julia> using DataTools, Transducers
+
+julia> firstitem(3:7)
+3
+
+julia> 3:7 |> Map(x -> x + 1) |> Filter(isodd) |> firstitem
+5
+```
+"""
+firstitem
+firstitem(xs::AbstractArray) = first(xs)
+firstitem(xs) = foldl(right, Take(1), xs)
+
+"""
+    lastitem(xs)
+
+Get the last item of `xs`.  Consume `xs` if necessary.
+
+# Examples
+```jldoctest
+julia> using DataTools, Transducers
+
+julia> lastitem(3:7)
+7
+
+julia> 3:7 |> Map(x -> x + 1) |> Filter(isodd) |> lastitem
+7
+```
+"""
+lastitem
+lastitem(xs::AbstractArray) = last(xs)
+lastitem(xs) = foldl(right, Map(identity), xs)
+
+"""
+    firstitems(xs, n::Integer)
+    firstitems(n::Integer) -> xs -> firstitems(xs, n)
+
+Get the first `n` items of `xs`.  Consume `xs` if necessary.
+"""
+firstitems
+firstitems(n::Integer) = xs -> firstitems(xs, n)
+firstitems(xs, n::Integer) = collect(Take(n), xs)
+firstitems(xs::AbstractArray, n::Integer) = view(xs, firstindex(xs):firstindex(xs)+n-1)
+
+"""
+    lastitems(xs, n::Integer)
+    lastitems(n::Integer) -> xs -> lastitems(xs, n)
+
+Get the last `n` items of `xs`.  Consume `xs` if necessary.
+"""
+lastitems
+lastitems(n::Integer) = xs -> lastitems(xs, n)
+lastitems(xs, n::Integer) = collect(TakeLast(n), xs)
+lastitems(xs::AbstractArray, n::Integer) = view(xs, lastindex(xs)-n+1:lastindex(xs))

test/test_firstitems.jl

@@ -0,0 +1,19 @@
+module TestFirstitems
+
+using DataTools
+using Test
+using Transducers
+
+include("utils.jl")
+
+@testset "firstitem" begin
+    @test firstitem(3:7) === 3
+    @test 3:7 |> Map(x -> x + 1) |> Filter(isodd) |> firstitem == 5
+end
+
+@testset "firstitems" begin
+    @test firstitems(3:7, 2) ==ₜ view(3:7, 1:2)
+    @test 3:7 |> Map(x -> x + 1) |> Filter(isodd) |> firstitems(2) == [5, 7]
+end
+
+end  # module

test/test_lastitems.jl

@@ -0,0 +1,19 @@
+module TestLastitems
+
+using DataTools
+using Test
+using Transducers
+
+include("utils.jl")
+
+@testset "lastitem" begin
+    @test lastitem(3:7) === 7
+    @test 3:7 |> Map(x -> x + 1) |> Filter(isodd) |> lastitem == 7
+end
+
+@testset "lastitems" begin
+    @test lastitems(3:7, 2) ==ₜ view(3:7, 4:5)
+    @test 3:7 |> Map(x -> x + 1) |> Filter(isodd) |> lastitems(2) == [5, 7]
+end
+
+end  # module

test/test_nitems.jl

@@ -0,0 +1,25 @@
+module TestNItems
+
+using DataTools
+using Test
+using Transducers
+using Transducers: IdentityTransducer
+
+@testset "_pop_innermost_maplikes" begin
+    pop(args...) = DataTools._pop_innermost_maplikes(opcompose(args...))
+    @test pop(Map(inv)) === IdentityTransducer()
+    @test pop(MapSplat(tuple), Map(inv)) === IdentityTransducer()
+    @test pop(Filter(isodd), MapSplat(tuple), Map(inv)) === Filter(isodd)
+    @test pop(Map(isodd), Filter(isodd), MapSplat(tuple), Map(inv)) ===
+          opcompose(Map(isodd), Filter(isodd))
+end
+
+@testset "nitems" begin
+    @test nitems(1:10) == 10
+    @test nitems(error(x) for x in 1:10) == 10
+    @test 1:10 |> Map(error) |> MapSplat(error) |> Scan(+) |> nitems == 10
+    @test 1:10 |> Filter(isodd) |> Map(error) |> MapSplat(error) |> nitems == 5
+    @test 1:10 |> Filter(isodd) |> Map(x -> x ÷ 3) |> Filter(isodd) |> nitems == 3
+end
+
+end  # module

test/utils.jl

@@ -0,0 +1,7 @@
+"""
+    ==ₜ(x, y)
+
+Check that _type_ and value of `x` and `y` are equal.
+"""
+==ₜ(_, _) = false
+==ₜ(x::T, y::T) where T = x == y