1*e1fe3e4aSElliott Hughesfrom fontTools.misc.visitor import Visitor 2*e1fe3e4aSElliott Hughesimport enum 3*e1fe3e4aSElliott Hughesimport pytest 4*e1fe3e4aSElliott Hughes 5*e1fe3e4aSElliott Hughes 6*e1fe3e4aSElliott Hughesclass E(enum.Enum): 7*e1fe3e4aSElliott Hughes E1 = 1 8*e1fe3e4aSElliott Hughes E2 = 2 9*e1fe3e4aSElliott Hughes E3 = 3 10*e1fe3e4aSElliott Hughes 11*e1fe3e4aSElliott Hughes 12*e1fe3e4aSElliott Hughesclass A: 13*e1fe3e4aSElliott Hughes def __init__(self): 14*e1fe3e4aSElliott Hughes self.a = 1 15*e1fe3e4aSElliott Hughes self.b = [2, 3] 16*e1fe3e4aSElliott Hughes self.c = {4: 5, 6: 7} 17*e1fe3e4aSElliott Hughes self._d = 8 18*e1fe3e4aSElliott Hughes self.e = E.E2 19*e1fe3e4aSElliott Hughes self.f = 10 20*e1fe3e4aSElliott Hughes 21*e1fe3e4aSElliott Hughes 22*e1fe3e4aSElliott Hughesclass B: 23*e1fe3e4aSElliott Hughes def __init__(self): 24*e1fe3e4aSElliott Hughes self.a = A() 25*e1fe3e4aSElliott Hughes 26*e1fe3e4aSElliott Hughes 27*e1fe3e4aSElliott Hughesclass TestVisitor(Visitor): 28*e1fe3e4aSElliott Hughes def __init__(self): 29*e1fe3e4aSElliott Hughes self.value = [] 30*e1fe3e4aSElliott Hughes 31*e1fe3e4aSElliott Hughes def _add(self, s): 32*e1fe3e4aSElliott Hughes self.value.append(s) 33*e1fe3e4aSElliott Hughes 34*e1fe3e4aSElliott Hughes def visitLeaf(self, obj): 35*e1fe3e4aSElliott Hughes self._add(obj) 36*e1fe3e4aSElliott Hughes super().visitLeaf(obj) 37*e1fe3e4aSElliott Hughes 38*e1fe3e4aSElliott Hughes 39*e1fe3e4aSElliott Hughes@TestVisitor.register(A) 40*e1fe3e4aSElliott Hughesdef visit(self, obj): 41*e1fe3e4aSElliott Hughes self._add("A") 42*e1fe3e4aSElliott Hughes 43*e1fe3e4aSElliott Hughes 44*e1fe3e4aSElliott Hughes@TestVisitor.register_attrs([(A, "e")]) 45*e1fe3e4aSElliott Hughesdef visit(self, obj, attr, value): 46*e1fe3e4aSElliott Hughes self._add(attr) 47*e1fe3e4aSElliott Hughes self._add(value) 48*e1fe3e4aSElliott Hughes return False 49*e1fe3e4aSElliott Hughes 50*e1fe3e4aSElliott Hughes 51*e1fe3e4aSElliott Hughes@TestVisitor.register(B) 52*e1fe3e4aSElliott Hughesdef visit(self, obj): 53*e1fe3e4aSElliott Hughes self._add("B") 54*e1fe3e4aSElliott Hughes self.visitObject(obj) 55*e1fe3e4aSElliott Hughes return False 56*e1fe3e4aSElliott Hughes 57*e1fe3e4aSElliott Hughes 58*e1fe3e4aSElliott Hughes@TestVisitor.register_attr(B, "a") 59*e1fe3e4aSElliott Hughesdef visit(self, obj, attr, value): 60*e1fe3e4aSElliott Hughes self._add("B a") 61*e1fe3e4aSElliott Hughes 62*e1fe3e4aSElliott Hughes 63*e1fe3e4aSElliott Hughesclass VisitorTest(object): 64*e1fe3e4aSElliott Hughes def test_visitor(self): 65*e1fe3e4aSElliott Hughes b = B() 66*e1fe3e4aSElliott Hughes visitor = TestVisitor() 67*e1fe3e4aSElliott Hughes visitor.visit(b) 68*e1fe3e4aSElliott Hughes assert visitor.value == ["B", "B a", "A", 1, 2, 3, 5, 7, "e", E.E2, 10] 69*e1fe3e4aSElliott Hughes 70*e1fe3e4aSElliott Hughes visitor.value = [] 71*e1fe3e4aSElliott Hughes visitor.defaultStop = True 72*e1fe3e4aSElliott Hughes visitor.visit(b) 73*e1fe3e4aSElliott Hughes assert visitor.value == ["B", "B a"] 74