diff --git a/stl/ast.py b/stl/ast.py index 0e80291..8a5f992 100644 --- a/stl/ast.py +++ b/stl/ast.py @@ -4,12 +4,9 @@ from collections import namedtuple, deque from itertools import repeat -from typing import Union from enum import Enum from sympy import Symbol -VarKind = Enum("VarKind", ["x", "u", "w"]) -str_to_varkind = {"x": VarKind.x, "u": VarKind.u, "w": VarKind.w} dt_sym = Symbol('dt', positive=True) t_sym = Symbol('t', positive=True) @@ -23,12 +20,7 @@ class AtomicPred(namedtuple("AP", ["id"])): class LinEq(namedtuple("LinEquality", ["terms", "op", "const"])): def __repr__(self): - n = len(self.terms) - rep = "{}" - if n > 1: - rep += " + {}"*(n - 1) - rep += " {op} {c}" - return rep.format(*self.terms, op=self.op, c=self.const) + return " + ".join(map(str, self.terms)) + f" {self.op} {self.const}" def children(self): return [] @@ -36,13 +28,12 @@ class LinEq(namedtuple("LinEquality", ["terms", "op", "const"])): class Var(namedtuple("Var", ["coeff", "id", "time"])): def __repr__(self): - time_str = "[{}]".format(self.time) - return "{c}*{i}{t}".format(c=self.coeff, i=self.id, t=time_str) + return f"{self.coeff}*{self.id}[{self.time}]" class Interval(namedtuple('I', ['lower', 'upper'])): def __repr__(self): - return "[{},{}]".format(self.lower, self.upper) + return f"[{self.lower},{self.upper}]" def children(self): return [self.lower, self.upper] @@ -51,14 +42,7 @@ class Interval(namedtuple('I', ['lower', 'upper'])): class NaryOpSTL(namedtuple('NaryOp', ['args'])): OP = "?" def __repr__(self): - n = len(self.args) - if n == 1: - return "{}".format(self.args[0]) - elif self.args: - rep = " {op} ".join(["({})"]*(len(self.args))) - return rep.format(*self.args, op=self.OP) - else: - return "" + return f" {self.OP} ".join(f"({x})" for x in self.args) def children(self): return self.args @@ -72,23 +56,23 @@ class And(NaryOpSTL): class ModalOp(namedtuple('ModalOp', ['interval', 'arg'])): + def __repr__(self): + return f"{self.OP}{self.interval}({self.arg})" + def children(self): return [self.arg] class F(ModalOp): - def __repr__(self): - return "◇{}({})".format(self.interval, self.arg) - + OP = "◇" class G(ModalOp): - def __repr__(self): - return "□{}({})".format(self.interval, self.arg) + OP = "□" class Neg(namedtuple('Neg', ['arg'])): def __repr__(self): - return "¬({})".format(self.arg) + return f"¬({self.arg})" def children(self): return [self.arg]