1  @doc "Python enlighten progressbar wrapper."
  2  module Pbar
  3  
  4  using Python
  5  using Python.PythonCall: pynew, keys, pyisnull, pycopy!
  6  using TimeTicks: now, Millisecond, Second
  7  
  8  const min_delta = Ref(Millisecond(0))
  9  const queued_counter = Ref(0)
 10  const enlighten = pynew()
 11  const emn = pynew()
 12  const pbar = pynew()
 13  const six = pynew()
 14  
 15  mutable struct PbarInstance
 16      pbar::Py
 17      fin::Bool
 18  end
 19  
 20  function clearpbar(pb)
 21      pb.pbar ∈ emn.counters.keys() && pbclose(pb)
 22  end
 23  
 24  function __init__()
 25      # Misc.__init__()
 26      @pymodule six
 27      @pymodule enlighten
 28      # pycopy!(enlighten, pyimport("enlighten"))
 29      pyisnull(emn) || emn.stop()
 30      pycopy!(emn, enlighten.get_manager())
 31      @debug @info "Pbar: Loaded enlighten."
 32  end
 33  
 34  @doc "Instantiate a progress bar:
 35  - `data`: `length(data)` determines the bar total
 36  - `unit`: what unit the display
 37  - `desc`: description will appear over the progressbar
 38  "
 39  macro pbar!(data, desc="", unit="", use_finalizer=false)
 40      data = esc(data)
 41      desc = esc(desc)
 42      unit = esc(unit)
 43      plu = esc(:pb_last_update)
 44      pb = esc(:pb)
 45      uf = esc(use_finalizer)
 46      quote
 47          @pbinit!
 48          !$pyisnull($pbar) && try
 49              $pbar.close(; clear=true)
 50          catch
 51          end
 52          pycopy!(
 53              $pbar, $emn.counter(; total=length($data), desc=$desc, unit=$unit, leave=false)
 54          )
 55          local $pb
 56          if $uf
 57              $pb = let $pb = PbarInstance($pbar, true)
 58                  finalizer($clearpbar, $pb)
 59              end
 60          else
 61              $pb = PbarInstance($pbar, false)
 62          end
 63          $min_delta[] = Millisecond($pyconvert(Float64, $pbar.min_delta) * 1e3)
 64          $pbar.refresh()
 65          local $plu = $now()
 66          $pb
 67      end
 68  end
 69  
 70  @doc "Single update to the progressbar with the new value."
 71  macro pbupdate!(n=1, args...)
 72      n = esc(n)
 73      plu = esc(:pb_last_update)
 74      pb = esc(:pb)
 75      quote
 76          let t = $now()
 77              if t - $plu > $min_delta[]
 78                  $pb.pbar.update($queued_counter[] + $n, $args...)
 79                  $plu = t
 80                  $queued_counter[] = 0
 81              else
 82                  $queued_counter[] += 1
 83              end
 84          end
 85      end
 86  end
 87  
 88  @doc "Terminates the progress bar."
 89  function pbclose(pb)
 90      pb.pbar.close(; clear=true)
 91  end
 92  
 93  macro pbclose()
 94      pb = esc(:pb)
 95      quote
 96          $pbclose($pb)
 97      end
 98  end
 99  
100  macro pbstop()
101      # FIXME: kwarg clear=true doesn't seem to work
102      quote
103          $pbar.close(; clear=true)
104          $emn.stop()
105      end
106  end
107  
108  macro pbinit!()
109      :($(__init__)())
110  end
111  
112  export @pbar!, @pbupdate!, @pbclose!, @pbstop, @pbinit!
113  
114  end