Fix count (#19)

JuliaFolds · Jul 5, 2020 · efb318f · efb318f
1 parent 9b1bf94
commit efb318f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/src/LazyGroupBy.jl b/src/LazyGroupBy.jl
@@ -150,12 +150,18 @@ impl(::typeof(prod), kwargs::NamedTuple, group::GroupedBy) =
 impl(::typeof(prod), kwargs::NamedTuple, f, group::GroupedBy) =
     impl(foldl, check_kwargs(prod, kwargs), mul_prod, Map(f), group)
 
+struct AsBool{F}
+    f::F
+end
+AsBool(::Type{T}) where {T} = AsBool{Type{T}}(T)
+@inline (f::AsBool)(x) = f.f(x)::Bool
+
 impl(::typeof(count), ::NamedTuple{()}, group::GroupedBy) =
-    impl(foldl, NamedTuple(), +, Map(identity), group)
+    impl(foldl, (init = 0,), +, Map(AsBool(identity)), group)
 impl(::typeof(count), ::NamedTuple{()}, f, group::GroupedBy) =
-    impl(foldl, NamedTuple(), +, Map(f), group)
+    impl(foldl, (init = 0,), +, Map(AsBool(f)), group)
 impl(::typeof(length), ::NamedTuple{()}, group::GroupedBy) =
-    impl(foldl, NamedTuple(), +, Map(_ -> 1), group)
+    impl(foldl, (init = 0,), +, Map(_ -> 1), group)
 
 # impl(::typeof(any), ::NamedTuple{()}, group::GroupedBy) =
 #     impl(foldl, NamedTuple(), |, IdentityTransducer(), group)

diff --git a/test/test_count.jl b/test/test_count.jl
@@ -8,6 +8,10 @@ using LazyGroupBy
           Dict(false => 3, true => 7)
     @test count.(grouped(identity, [true, false, true, true, true])) ==
           Dict(false => 0, true => 4)
+    @test_throws(
+        Exception,
+        count.(_ -> 0.5, grouped(isodd, [0, 7, 3, 1, 5, 9, 4, 3, 0, 5])),
+    )
 end
 
 end  # module