calc2.py

#!/usr/bin/env python
#
# by Erik Osheim

from pymeta.grammar import OMeta
from itertools import chain
import math
import sys

# some errors we can encounter during compilation and execution
class VarError(Exception): pass
class FunctionError(Exception): pass
class OperatorError(Exception): pass

# ast node to represent a literal number, e.g. "42" or "3.14"
class Num(object):
    def __init__(self, n):
        self.value = n
        self.vars = set()
    def __str__(self):
        return str(self.value)
    def run(self, kw):
        return self.value

# ast node to represent a single variable number, e.g. "qux13" or "FooBar"
class Var(object):
    def __init__(self, s):
        self.name = s
        self.vars = set([s])
    def __str__(self):
        return self.name
    def run(self, kw):
        if self.name in kw:
            return kw[self.name]
        else:
            raise VarError("undefined variable %r" % self.name)

# ast node to represent function application, e.g. "cos(x * pi)"
class Apply(object):
    fs = {
        'abs': lambda xs: abs(xs[0]),
        'ceil': lambda xs: math.ceil(xs[0]),
        'cos': lambda xs: math.cos(xs[0]),
        'degrees': lambda xs: math.degrees(xs[0]),
        'floor': lambda xs: math.floor(xs[0]),
        'log': lambda xs: math.log(xs[0]),
        'log10': lambda xs: math.log10(xs[0]),
        'log2': lambda xs: math.log(xs[0], 2),
        'max': lambda xs: max(xs),
        'min': lambda xs: min(xs),
        'radians': lambda xs: math.radians(xs[0]),
        'round': lambda xs: round(xs[0]),
        'sin': lambda xs: math.sin(xs[0]),
        'tan': lambda xs: math.tan(xs[0]),
    }
    def __init__(self, s, args):
        self.name = s
        self.args = args
        self.vars = reduce(lambda a, b: a | b, [a.vars for a in args])
        self.f = self.fs.get(s)
        if self.f is None:
            raise FunctionError("undefined function %r" % s)
    def __str__(self):
        return '(%s %s)' % (self.name, ' '.join([str(a) for a in self.args]))
    def run(self, kw):
        return self.f([a.run(kw) for a in self.args])

# ast node to represnt a unary operator, e.g. "not x", or "9!"
class Unop(object):
    fs = {
        '-': lambda x: -x,
        '!': lambda x: math.factorial(x),
        'not': lambda x: int(not x),
    }
    def __init__(self, op, arg):
        self.op = op
        self.arg = arg
        self.vars = arg.vars
        self.f = self.fs.get(op)
        if self.f is None:
            raise OperatorError("undefined operator %r" % op)
    def __str__(self):
        return '(%s %s)' % (self.op, self.arg)
    def run(self, kw):
        return self.f(self.arg.run(kw))

# ast node to represent a binary operator, e.g. "2 + 2" or "x and y"
class Binop(object):
    fs = {
        '+': lambda x, y: x + y,
        '-': lambda x, y: x - y,
        '*': lambda x, y: x * y,
        '/': lambda x, y: x / y,
        '//': lambda x, y: x // y,
        '%': lambda x, y: x % y,
        #'^': lambda x, y: x ** y,
        '**': lambda x, y: x ** y,

        '<<': lambda x, y: x << y,
        '>>': lambda x, y: x >> y,
        '&': lambda x, y: x & y,
        '|': lambda x, y: x | y,
        '^': lambda x, y: x ^ y,

        '==': lambda x, y: int(x == y),
        '!=': lambda x, y: int(x != y),
        '<': lambda x, y: int(x < y),
        '<': lambda x, y: int(x <= y),
        '>=': lambda x, y: int(x >= y),
        '>': lambda x, y: int(x > y),
        '<=>': lambda x, y: cmp(x, y),

        'and': lambda x, y: int(x and y),
        'or': lambda x, y: int(x or y),
    }
    def __init__(self, op, lhs, rhs):
        self.op = op
        self.lhs = lhs
        self.rhs = rhs
        self.vars = lhs.vars | rhs.vars
        self.f = self.fs.get(op)
        if self.f is None:
            raise OperatorError("undefined operator %r" % op)
    def __str__(self):
        return '(%s %s %s)' % (self.op, self.lhs, self.rhs)
    def run(self, kw):
        return self.f(self.lhs.run(kw), self.rhs.run(kw))

# ast node to represent a ternary expression, e.g. if x > 0 then x else y
class Ternary(object):
    def __init__(self, c, t, f):
        self.condition = c
        self.true = t
        self.false = f
        self.vars = c.vars | t.vars | f.vars
    def __str__(self):
        return '(if %s then %s else %s)' % (self.condition, self.true, self.false)
    def run(self, kw):
        if self.condition.run(kw):
            return self.true.run(kw)
        else:
            return self.false.run(kw)        

rules = """
# top-level program production. this is the full expression we're parsing.
prog ::= <expr>

# expression-level productions. each of these deals with a particular type
# of expression. the different rules are necessary for operator precedence.
# naming convention "expr_xyz" means "deals with xyz operators".
expr ::= <expr_cond>

expr_cond ::= <if> <s> <expr>:c <s> <then> <s> <expr>:t <s> <endif>:f => Ternary(c, t, f)
            | <expr_or>:e => e

endif ::= <elif> <s> <expr>:c <s> <then> <s> <expr>:t <s> <endif>:f => Ternary(c, t, f)
        | <else> <s> <expr>:e => e

expr_or ::= <expr_or>:lhs <s> <or> <s> <expr_and>:rhs => Binop('or', lhs, rhs)
          | <expr_and>:e => e

expr_and ::= <expr_and>:lhs <s> <and> <s> <expr_not>:rhs => Binop('and', lhs, rhs)
           | <expr_not>:e => e

expr_not ::= <not> <s> <expr_not>:e => Unop('not', e)
           | <expr_cmp>:e => e

expr_cmp ::= <expr_cmp>:lhs <s> <compare>:op <s> <expr_shift>:rhs => Binop(op, lhs, rhs)
           | <expr_shift>:e => e

expr_shift ::= <expr_shift>:lhs <s> <shift>:op <s> <expr_add>:rhs => Binop(op, lhs, rhs)
             | <expr_add>:e => e

expr_add ::= <expr_add>:lhs <s> '+' <s> <expr_mul>:rhs => Binop('+', lhs, rhs)
           | <expr_add>:lhs <s> '-' <s> <expr_mul>:rhs => Binop('-', lhs, rhs)
           | <expr_mul>:e => e

expr_mul ::= <expr_mul>:lhs <s> '*' <s> <expr_neg>:rhs => Binop('*', lhs, rhs)
           | <expr_mul>:lhs <s> <intdiv> <s> <expr_neg>:rhs => Binop('//', lhs, rhs)
           | <expr_mul>:lhs <s> '/' <s> <expr_neg>:rhs => Binop('/', lhs, rhs)
           | <expr_mul>:lhs <s> '%' <s> <expr_neg>:rhs => Binop('%', lhs, rhs)
           | <expr_neg>:e => e

expr_neg ::= '-' <s> <expr_neg>:e => Unop('-', e)
           | <expr_exp>:e => e

expr_exp ::= <expr_fact>:lhs <s> <pow> <s> <expr_neg>:rhs => Binop('**', lhs, rhs)
           | <expr_fact>:e => e

expr_fact ::= <expr_fact>:e <s> '!' => Unop('!', e)
            | <expr_num>:e => e

expr_num ::= <apply>:a <s> => a
           | <id>:s <s> => Var(s)
           | <num>:n <s> => Num(n)
           | '(' <s> <expr>:e <s> ')' <s> => e
           | '|' <s> <expr>:e <s> '|' <s> => Unop('abs', e)

# low-level atomic numbers. these are literal values, variables, and function
# application.
apply ::= <id>:name <s> '(' <s> <args>:ns ')' <s> => Apply(name, ns)

args ::= <expr>:e <s> ',' <s> <args>:es => [e] + es
       | <expr>:e => [e]

id  ::= <letter>:c <letterOrDigit>*:cs => c + ''.join(cs)

num ::= <dec>:d => d
      | <int>:n => n

# these productions represent terminals in the grammar, like literal strings,
# numbers, etc.
s ::= <spaces>

or ::= 'o' 'r'
and ::= 'a' 'n' 'd'
not ::= 'n' 'o' 't'
if ::= 'i' 'f'
then ::= 't' 'h' 'e' 'n'
elif ::= 'e' 'l' 'i' 'f'
else ::= 'e' 'l' 's' 'e'

compare ::= '<' '=' '>' => '<=>'
          | '<' '=' => '<='
          | '<' => '<'
          | '>' '=' => '>='
          | '>' => '>'
          | '!' '=' => '!='
          | '=' '=' => '='

shift ::= '<' '<' => '<<'
        | '>' '>' => '>>'

intdiv ::= '/' '/'
pow ::= '*' '*'

zero    ::= '0'
nonzero ::= '1' | '2' | '3' | '4' | '5' | '6' | '7' | '8' | '9'
digit   ::= <zero> | <nonzero>

int ::= <nonzero>:c <digit>*:cs => int(c + ''.join(cs))
      | <zero> => 0

dec ::= <int>:n '.' <digit>+:cs => float(str(n) + '.' + ''.join(cs))
"""

def parsearg(s):
    k, v = s.split('=', 1)
    if '.' in v:
        return (k, float(v))
    else:
        return (k, int(v))

def main():
    # if no arguments provided, show usage
    if len(sys.argv) < 2:
        print "usage: %s expr [var1=number var2=number ...]" % sys.argv[0]
        return

    # build the grammar based on our string of rules
    g = OMeta.makeGrammar(rules, globals(), name='calc')
    
    # get the prog we want to execute
    prog = sys.argv[1]
    
    # apply our grammar, starting with the "prog" non-terminal
    r, e = g(prog).apply("prog")
    
    # see if we got a result, and if we parsed all our input
    if r is None or e[0] != len(prog):
        print "parse error near %r (column %d)" % (prog[e[0]], e[0])
        return

    # print a lispy AST representation for the user, along with what
    # variables (if any) the expression requires.
    print 'parsed: %s' % r
    if r.vars:
        print 'variables: %s' % ' '.join(r.vars)
    else:
        print 'no variables'

    # parse the provided variables, and evaluate the expression
    try:
        kw = dict([parsearg(s) for s in sys.argv[2:]])
        n = r.run(kw)
        print "result: %r" % n
    except Exception, e:
        print "error: %s" % e

if __name__ == "__main__":
    main()