[aggr-] allow mean and stdev for dates and timedeltas

visidata/aggregators.py

@@ -3,6 +3,7 @@
 import functools
 import collections
 import statistics
+import datetime
 
 from visidata import Progress, Sheet, Column, ColumnsSheet, VisiData
 from visidata import vd, anytype, vlen, asyncthread, wrapply, AttrDict, date
@@ -107,14 +108,34 @@ def _funcRows(col, rows): # wrap builtins so they can have a .type
 def mean(vals):
  vals = list(vals)
  if vals:
- return float(sum(vals))/len(vals)
+ if type(vals[0]) is date:
+ vals = [d.timestamp() for d in vals]
+ ans = float(sum(vals))/len(vals)
+ return datetime.date.fromtimestamp(ans)
+ elif isinstance(vals[0], datetime.timedelta):
+ return datetime.timedelta(seconds=vsum(vals)/datetime.timedelta(seconds=len(vals)))
+ else:
+ return float(sum(vals))/len(vals)
 
 def _vsum(vals):
  return sum(vals, start=type(vals[0] if len(vals) else 0)()) #1996
 
 # start parameter in sum() added in Python 3.8
 vsum = _vsum if sys.version_info[:2] >= (3, 8) else sum
 
+def stdev(vals):
+ if vals and len(vals) >= 2:
+ if type(vals[0]) is date:
+ vals = [d.timestamp() for d in vals]
+ return datetime.timedelta(seconds=statistics.stdev(vals))
+ elif isinstance(vals[0], datetime.timedelta):
+ vals = [d.total_seconds() for d in vals]
+ return datetime.timedelta(seconds=statistics.stdev(vals))
+ return statistics.stdev(vals)
+ else:
+ vd.error('stdev requires at least two data points')
+ return None
+
 # http://code.activestate.com/recipes/511478-finding-the-percentile-of-the-values/
 def _percentile(N, percent, key=lambda x:x):
  """
@@ -148,15 +169,15 @@ def quantiles(q, helpstr):
 
 vd.aggregator('min', min, 'minimum value')
 vd.aggregator('max', max, 'maximum value')
-vd.aggregator('avg', mean, 'arithmetic mean of values', type=float)
-vd.aggregator('mean', mean, 'arithmetic mean of values', type=float)
+vd.aggregator('avg', mean, 'arithmetic mean of values', type=anytype)
+vd.aggregator('mean', mean, 'arithmetic mean of values', type=anytype)
 vd.aggregator('median', statistics.median, 'median of values')
 vd.aggregator('mode', statistics.mode, 'mode of values')
 vd.aggregator('sum', vsum, 'sum of values')
 vd.aggregator('distinct', set, 'distinct values', type=vlen)
 vd.aggregator('count', lambda values: sum(1 for v in values), 'number of values', type=int)
 vd.aggregator('list', list, 'list of values', type=anytype)
-vd.aggregator('stdev', statistics.stdev, 'standard deviation of values', type=float)
+vd.aggregator('stdev', stdev, 'standard deviation of values', type=anytype)
 
 vd.aggregators['q3'] = quantiles(3, 'tertiles (33/66th pctile)')
 vd.aggregators['q4'] = quantiles(4, 'quartiles (25/50/75th pctile)')
@@ -225,7 +246,11 @@ def memo_aggregate(col, agg_choices, rows):
  for agg in aggs:
  aggval = agg(col, rows)
  typedval = wrapply(agg.type or col.type, aggval)
- dispval = col.format(typedval)
+ if agg.name == 'stdev' and (col.type is date):
+ # col type is a date, but typedval is a timedelta
+ dispval = str(typedval)
+ else:
+ dispval = col.format(typedval)
  k = col.name+'_'+agg.name
  vd.status(f'{k}={dispval}')
  vd.memory[k] = typedval