Shunting-yard algorithm in Python
Converts infix expressions to postfix or builds expression trees using an operator stack.
- File
shunting_yard_algorithm.py, 87 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/shunting-yard-algorithm/. The authors are listed in prov/2TMV07.jsonl.
"""Dijkstra's shunting-yard algorithm: infix tokens to postfix tokens.
A token is a string. "(" and ")" are parentheses, a key of the operator
table is a binary operator, and any other token is an operand. The operator
table maps each operator to (precedence, associativity), where a larger
precedence binds tighter and associativity is "left" or "right". All
operators of one precedence level must have the same associativity.
"""
from typing import Mapping, Sequence
DEFAULT_OPS: dict[str, tuple[int, str]] = {
"+": (1, "left"),
"-": (1, "left"),
"*": (2, "left"),
"/": (2, "left"),
"^": (3, "right"),
}
def _pops(top: tuple[int, str], t: tuple[int, str]) -> bool:
"""True when the operator on the stack must be output before t is pushed."""
return top[0] > t[0] or (top[0] == t[0] and t[1] == "left")
def shunting_yard(tokens: Sequence[str], ops: Mapping[str, tuple[int, str]] = DEFAULT_OPS) -> list[str]:
"""Return the postfix form of the infix expression `tokens`.
Raises ValueError when the input is not a well-formed infix expression
(unbalanced parentheses, two operands or two operators in a row, an
operator at the start or the end, or an empty expression).
Time and extra space Theta(n) for n tokens: each token is pushed and
popped at most once.
"""
out: list[str] = []
stack: list[str] = [] # @line L1
expect_operand = True # the input check: operands and operators alternate
for t in tokens: # @line L2
if t == "(":
if not expect_operand:
raise ValueError("an operand or '(' cannot follow an operand")
stack.append(t) # @line L4
elif t == ")": # @line L5
if expect_operand:
raise ValueError("')' must follow an operand")
while stack and stack[-1] != "(": # @line L6
out.append(stack.pop())
if not stack:
raise ValueError("unbalanced ')'")
stack.pop() # @line L7
elif t in ops: # @line L8
if expect_operand:
raise ValueError(f"operator {t!r} needs a left operand")
p = ops[t]
while stack and stack[-1] != "(" and _pops(ops[stack[-1]], p): # @line L9
out.append(stack.pop()) # @line L10
stack.append(t) # @line L11
expect_operand = True
else:
if not expect_operand:
raise ValueError("two operands in a row")
out.append(t) # @line L3
expect_operand = False
if expect_operand:
raise ValueError("the expression is empty or ends with an operator")
while stack: # @line L12
if stack[-1] == "(": # @line L13
raise ValueError("unbalanced '('")
out.append(stack.pop()) # @line L14
return out # @line L15
def evaluate_postfix(postfix: Sequence[str], ops: Mapping[str, tuple[int, str]] = DEFAULT_OPS) -> object:
"""Build the expression tree of a postfix sequence as nested tuples (op, left, right)."""
st: list[object] = []
for t in postfix:
if t in ops:
b, a = st.pop(), st.pop()
st.append((t, a, b))
else:
st.append(t)
if len(st) != 1:
raise ValueError("not a postfix expression")
return st[0]
test_shunting_yard_algorithm.py 185 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/shunting-yard-algorithm/. The authors are listed in prov/2TMV07.jsonl.
import itertools
import random
import pytest
from shunting_yard_algorithm import evaluate_postfix, shunting_yard
# The unambiguous precedence grammar that defines the meaning of an expression.
GRAMMAR = [
("E", ("E", "+", "T")), ("E", ("E", "-", "T")), ("E", ("T",)),
("T", ("T", "*", "F")), ("T", ("T", "/", "F")), ("T", ("F",)),
("F", ("P", "^", "F")), ("F", ("P",)),
("P", ("(", "E", ")")), ("P", ("a",)), ("P", ("b",)),
]
NTS = {"E", "T", "F", "P"}
def spans(w):
"""Least fixpoint of the facts 'A derives w[i:j]' (an independent oracle)."""
n = len(w)
known = set()
def ends_of(rhs, i, j):
ends = {i}
for x in rhs:
nxt = set()
for p in ends:
if x in NTS:
nxt |= {q for q in range(p, j + 1) if (x, p, q) in known}
elif p < j and w[p] == x:
nxt.add(p + 1)
ends = nxt
return ends
changed = True
while changed:
changed = False
for a, rhs in GRAMMAR:
for i in range(n + 1):
for j in range(i, n + 1):
if (a, i, j) not in known and j in ends_of(rhs, i, j):
known.add((a, i, j))
changed = True
return known
def grammar_tree(w):
"""The unique parse tree of w, as nested (op, left, right) tuples, or None."""
known = spans(w)
n = len(w)
if ("E", 0, n) not in known:
return None
def build(a, i, j):
for lhs, rhs in GRAMMAR:
if lhs != a:
continue
for cut in itertools.product(range(i, j + 1), repeat=len(rhs) - 1):
bounds = (i,) + cut + (j,)
if any(bounds[k] > bounds[k + 1] for k in range(len(rhs))):
continue
ok = True
for x, p, q in zip(rhs, bounds, bounds[1:]):
if x in NTS:
ok = ok and (x, p, q) in known
else:
ok = ok and q == p + 1 and w[p] == x
if not ok:
continue
if len(rhs) == 3 and rhs[1] in "+-*/^":
return (rhs[1], build(rhs[0], bounds[0], bounds[1]), build(rhs[2], bounds[2], bounds[3]))
if rhs[0] == "(":
return build("E", bounds[1], bounds[2])
if rhs[0] in NTS:
return build(rhs[0], i, j)
return w[i]
raise AssertionError("no split")
return build("E", 0, n)
def test_example() -> None:
toks = "a + b * ( c - d ) ^ e ^ f - g".split()
assert shunting_yard(toks) == "a b c d - e f ^ ^ * + g -".split()
def test_associativity_and_precedence() -> None:
assert shunting_yard("a - b - c".split()) == "a b - c -".split()
assert shunting_yard("a ^ b ^ c".split()) == "a b c ^ ^".split()
assert shunting_yard("a + b * c".split()) == "a b c * +".split()
assert shunting_yard("( a + b ) * c".split()) == "a b + c *".split()
assert shunting_yard(["x"]) == ["x"]
assert shunting_yard("( ( x ) )".split()) == ["x"]
@pytest.mark.parametrize("bad", ["", "+", "a +", "+ a", "a b", "( a", "a )", "( )", ") a (", "a ( b )", "a + * b", "( a + ) b"])
def test_rejects_malformed(bad: str) -> None:
with pytest.raises(ValueError):
shunting_yard(bad.split())
def random_tree(rng, depth):
if depth == 0 or rng.random() < 0.3:
return rng.choice("abcdxyz")
return (rng.choice("+-*/^"), random_tree(rng, depth - 1), random_tree(rng, depth - 1))
PREC = {"+": 1, "-": 1, "*": 2, "/": 2, "^": 3}
def infix(tree, rng):
"""Print a tree with the parentheses it needs, plus some extra ones."""
if isinstance(tree, str):
s = [tree]
else:
op, l, r = tree
ls, rs = infix(l, rng), infix(r, rng)
lneed = not isinstance(l, str) and (PREC[l[0]] < PREC[op] or (PREC[l[0]] == PREC[op] and op == "^"))
rneed = not isinstance(r, str) and (PREC[r[0]] < PREC[op] or (PREC[r[0]] == PREC[op] and op != "^"))
s = (["("] + ls + [")"] if lneed else ls) + [op] + (["("] + rs + [")"] if rneed else rs)
if rng.random() < 0.1:
s = ["("] + s + [")"]
return s
def postorder(tree):
out, stack = [], [(tree, False)]
while stack:
t, done = stack.pop()
if isinstance(t, str):
out.append(t)
elif done:
out.append(t[0])
else:
stack.append((t, True))
stack.append((t[2], False))
stack.append((t[1], False))
return out
def test_random_trees() -> None:
rng = random.Random(1961)
for _ in range(1500):
tree = random_tree(rng, rng.randint(0, 5))
toks = infix(tree, rng)
post = shunting_yard(toks)
assert post == postorder(tree)
assert evaluate_postfix(post) == tree
def test_random_token_strings_against_grammar() -> None:
rng = random.Random(42)
alphabet = ["a", "b", "+", "-", "*", "^", "(", ")"]
for _ in range(1500):
w = [rng.choice(alphabet) for _ in range(rng.randint(0, 9))]
expected = grammar_tree(w)
if expected is None:
with pytest.raises(ValueError):
shunting_yard(w)
else:
assert evaluate_postfix(shunting_yard(w)) == expected
def test_all_short_strings() -> None:
alphabet = ["a", "-", "^", "(", ")"]
for n in range(0, 6):
for w in itertools.product(alphabet, repeat=n):
w = list(w)
expected = grammar_tree(w)
if expected is None:
with pytest.raises(ValueError):
shunting_yard(w)
else:
assert evaluate_postfix(shunting_yard(w)) == expected
def test_long_input() -> None:
n = 20000
toks = ["("] * n + ["a"] + [")"] * n
assert shunting_yard(toks) == ["a"]
toks = ["a"] + ["^", "a"] * n
assert shunting_yard(toks) == ["a"] * (n + 1) + ["^"] * n
What the next level needs
- the tests did not run in this build