文字列の数式を計算するやつ

はじめに

ひょんなことから文字列の数式（四則演算）を計算することになり、Pythonで実装したので載せておきます。（floatは計算できません！！）

コード

import os
import re

ZERO = ord("0")
NINE = ord("9")

def paren(st):
    if st[0] == "(":
        ans, idx = first(st[1:])
        return ans, idx+2
    elif ZERO <= ord(st[0]) <= NINE:
        i = 1
        while i < len(st) and ZERO <= ord(st[i]) <= NINE:
            i += 1
        return int(st[:i]), i
    return 0, 0

def second(st):
    ans, idx = paren(st)

    i = idx
    while i < len(st):
        if st[i] == "*":
            tmp, idx = paren(st[i+1:])
            ans *= tmp
            i += idx+1
        elif st[i] == "/":
            tmp, idx = paren(st[i+1:])
            ans /= tmp
            i += idx+1
        else:
            return ans, i
    return ans, i

def first(st):
    ans, idx = second(st)

    i = idx
    while i < len(st):
        if st[i] == "+":
            tmp, idx = second(st[i+1:])
            ans += tmp
            i += idx+1
        elif st[i] == "-":
            tmp, idx = second(st[i+1:])
            ans -= tmp
            i += idx+1
        else:
            return ans, i
    return ans, i

def calc(s):
    if s.count("(") != s.count(")") or re.search("[^\+\-\*\/()0-9]", s):
        return "ERROR"
    else:
        return first(s)[0]

数式において、優先順位は
括弧　＞　掛け算割り算　＞　足し算割り算
です。
それぞれの処理を順に paren、second、first 関数にまかせています。

calc(s)

最初に数式が入る関数。
括弧の数が合わなかったり、四則演算に相応しくない文字が含まれている場合はERRORを返し、適切な文字列である場合はfirst関数に渡します。

paren(st)

first() -> second() と入力がたらい回しにされて一番最初に処理が行われる関数。
”(”が最初にある場合は”(”の１つ先以降をfirstに渡します（新たな式として処理）。
数字が最初にある場合は数字が続くまでインデックスを進め、intに変換します。
返り値は
(計算値,進んだインデックス)
になります。

second(st)

parenの次に処理される関数
parenで得られた値ansに対し、次にある値を掛けるか割るかします。
返り値はparenと同様です。

first(st)

secondの次に処理される関数
上記同様、こちらは足すか引くかします。
返り値も同様です。

最後に

インデックスの管理が大変だった・・・
"+("や"-("を直接paren内でやろうとしてめっちゃ詰まった・・・
（firstで+,-の後をsecondに渡せば完了）

参考

数式文字列の四則演算(整数値) - プログラミングのメモ帳