Wrap result of tracktime in a struct to have a custom show method

and avoid unconditional printing of timing results

Author: andreasnoack

src/diagnostics.jl

@@ -84,6 +84,12 @@ function stop_tracking_time()
     compute_context[] = nothing
 end
 
+struct TrackTime
+    obj
+end
+
+Base.show(io::IO, ::MIME"text/plain", t::TrackTime) = show_timings(t.obj, maxdepth=5)
+
 """
 `tracktime(f)`
 
@@ -96,8 +102,7 @@ function tracktime(f; profile=false, maxdepth=5)
     ctx = compute_context[]
     stop_tracking_time()
     t = fetch_timings!(ctx, profile=profile)
-    show_timings(t, maxdepth=maxdepth)
-    t, res
+    TrackTime(t), res
 end
 
 function fetch_timings!(ctx=get_context(); profile=true)

src/ml.jl

@@ -25,11 +25,6 @@ function schema(xs::DArray)
                        delayedmap(x -> schema(x), xs.chunks)))
 end
 
-function schema(xs::DArray, T)
-    collect(treereduce(delayed(merge),
-                       delayedmap(x -> schema(x, T), xs.chunks)))
-end
-
 schema(xs::ArrayOp) = schema(compute(xs))
 schema(xs::ArrayOp, T) = schema(compute(xs), T)
 
@@ -87,6 +82,13 @@ function schema(xs::PooledArray)
     schema(xs, Categorical)
 end
 
+for T in (:Continuous, :Categorical)
+    @eval function schema(xs::DArray, ::Type{$T})
+        collect(treereduce(delayed(merge),
+                       delayedmap(x -> schema(x, $T), xs.chunks)))
+    end
+end
+
 Base.keys(c::Categorical) = keys(c.series.stats[1])
 width(c::Categorical) = length(keys(c))
 merge(c1::Categorical, c2::Categorical) = Categorical(merge(c1.series, c2.series))