added var_names property

This commit is contained in:
Marcell Vazquez-Chanlatte 2017-11-26 20:01:54 -08:00
parent 034bc22955
commit 9098cff929
2 changed files with 13 additions and 1 deletions

View file

@ -5,7 +5,7 @@ from functools import lru_cache
import funcy as fn import funcy as fn
import lenses import lenses
from lenses import lens from lenses import lens, bind
def flatten_binary(phi, op, dropT, shortT): def flatten_binary(phi, op, dropT, shortT):
@ -77,6 +77,12 @@ class AST(object):
def atomic_predicates(self): def atomic_predicates(self):
return set(AP_lens(self).Each().collect()) return set(AP_lens(self).Each().collect())
@property
def var_names(self):
symbols = set(bind(self.lineqs).Each().terms.Each().collect())
symbols |= self.atomic_predicates
return set(bind(symbols).Each().id.collect())
def inline_context(self, context): def inline_context(self, context):
phi, phi2 = self, None phi, phi2 = self, None

View file

@ -45,3 +45,9 @@ def test_walk():
phi = stl.parse( phi = stl.parse(
'((G[0, 1](x + y > a?)) & (F[1,2](z - x > 0))) | ((X(AP1)) U (AP2))') '((G[0, 1](x + y > a?)) & (F[1,2](z - x > 0))) | ((X(AP1)) U (AP2))')
assert len(list((~phi).walk())) == 11 assert len(list((~phi).walk())) == 11
def test_var_names():
phi = stl.parse(
'((G[0, 1](x + y > a?)) & (F[1,2](z - x > 0))) | ((X(AP1)) U (AP2))')
assert phi.var_names == {'x', 'y', 'z', 'x', 'AP1', 'AP2'}