diff --git a/scripts/csv.py b/scripts/csv.py index 21e34ad1..f6bc515e 100755 --- a/scripts/csv.py +++ b/scripts/csv.py @@ -654,7 +654,12 @@ class CsvExpr: # expr node base class class Expr: - def __init__(self, *args): + def __init__(self, *args, by=None): + self.by = by or [] + for k in self.by: + if not isinstance(k, CsvExpr.Field): + raise CsvExpr.Error("complicated by? %s" % k) + for k, v in zip('abcdefghijklmnopqrstuvwxyz', args): setattr(self, k, v) @@ -668,8 +673,10 @@ class CsvExpr: return sum(1 for _ in self) def __repr__(self): - return '%s(%s)' % ( + return '%s%s(%s)' % ( self.__class__.__name__, + '[%s]' % (','.join(repr(k) for k in self.by)) + if self.by else '', ','.join(repr(v) for v in self)) def fields(self): @@ -741,9 +748,14 @@ class CsvExpr: # func expr helper def func(funcs): - def func(name, args="a"): + def func(name, by=None, args=None): + if by is None and args is None: + by, args = None, "a" + elif args is None: + by, args = None, by def func(f): f._func = name + f._fby = by f._fargs = args funcs[f._func] = f return f @@ -960,22 +972,16 @@ class CsvExpr: return CsvGStddev()([v.eval(fields, state) for v in self]) # enumerate exprs - @func('enumerate', '[*by]') + @func('enumerate', 'by', '') class Enumerate(Expr): """A [per by] number incremented for each result""" def fields(self): - # don't typecheck by fields return set() def type(self, types={}): - # don't typecheck, but make sure we can read by fields - for v in self: - if not isinstance(v, CsvExpr.Field): - raise CsvExpr.Error("complicated by field? %s" % v) return CsvInt def fold(self, types={}): - # don't typecheck by fields return CsvSum def eval(self, fields={}, state=None): @@ -984,7 +990,7 @@ class CsvExpr: # enumerate k = ('enumerate', id(self)) + tuple( - fields.get(v.a) for v in self) + fields.get(k.a) for k in self.by) x = state.get(k) if x is None: z = 0 @@ -994,25 +1000,9 @@ class CsvExpr: state[k] = z return CsvInt(z) - @func('accumulate', 'a[, *by]') + @func('accumulate', 'by', 'a') class Accumulate(Expr): """A [per by] running sum across results""" - def fields(self): - # don't typecheck by fields - return self.a.fields() - - def type(self, types={}): - # don't typecheck, but make sure we can read by fields - t = self.a.type(types) - for v in it.islice(self, 1, None): - if not isinstance(v, CsvExpr.Field): - raise CsvExpr.Error("complicated by field? %s" % v) - return t - - def fold(self, types={}): - # don't typecheck by fields - return self.a.fold(types) - def eval(self, fields={}, state=None): y = self.a.eval(fields, state) if state is None: @@ -1020,7 +1010,7 @@ class CsvExpr: # accumulate k = ('accumulate', id(self)) + tuple( - fields.get(v.a) for v in it.islice(self, 1, None)) + fields.get(k.a) for k in self.by) x = state.get(k) if x is None: z = y @@ -1030,25 +1020,9 @@ class CsvExpr: state[k] = z return z - @func('delta', 'a[, *by]') + @func('delta', 'by', 'a') class Delta(Expr): """A [per by] difference between subsequent results""" - def fields(self): - # don't typecheck by fields - return self.a.fields() - - def type(self, types={}): - # don't typecheck, but make sure we can read by fields - t = self.a.type(types) - for v in it.islice(self, 1, None): - if not isinstance(v, CsvExpr.Field): - raise CsvExpr.Error("complicated by field? %s" % v) - return t - - def fold(self, types={}): - # don't typecheck by fields - return self.a.fold(types) - def eval(self, fields={}, state=None): y = self.a.eval(fields, state) if state is None: @@ -1056,7 +1030,7 @@ class CsvExpr: # compute delta k = ('delta', id(self)) + tuple( - fields.get(v.a) for v in it.islice(self, 1, None)) + fields.get(k.a) for k in self.by) x = state.get(k) if x is None: z = y @@ -1531,7 +1505,12 @@ class CsvExpr: print('funcs:') for func in cls.funcs.keys(): print(' %-21s %s' % ( - '%s(%s)' % (func, CsvExpr.funcs[func]._fargs), + '%s%s(%s)' % ( + func, + '[%s]' % CsvExpr.funcs[func]._fby + if CsvExpr.funcs[func]._fby + else '', + CsvExpr.funcs[func]._fargs), CsvExpr.funcs[func].__doc__)) # parse an expr @@ -1560,6 +1539,25 @@ class CsvExpr: elif p.match('[_a-zA-Z][_a-zA-Z0-9]*'): a = p.chomp() + by = None + if p.match('\['): + p.chomp() + if a not in CsvExpr.funcs: + raise CsvExpr.Error("unknown function? %s" % a) + by = [] + while True: + if not p.match('\]'): + b = p_expr(p) + by.append(b) + if p.match(','): + p.chomp() + continue + if not p.match('\]'): + raise CsvExpr.Error("mismatched squares? %s" % p) + p.chomp() + break + + args = None if p.match('\('): p.chomp() if a not in CsvExpr.funcs: @@ -1575,8 +1573,12 @@ class CsvExpr: if not p.match('\)'): raise CsvExpr.Error("mismatched parens? %s" % p) p.chomp() - a = CsvExpr.funcs[a](*args) break + + if args is not None: + a = CsvExpr.funcs[a](*args, by=by) + elif by is not None: + raise CsvExpr.Error("expected parens? %s" % p) else: a = CsvExpr.Field(a)