1*e1fe3e4aSElliott Hughes"""Generic visitor pattern implementation for Python objects.""" 2*e1fe3e4aSElliott Hughes 3*e1fe3e4aSElliott Hughesimport enum 4*e1fe3e4aSElliott Hughes 5*e1fe3e4aSElliott Hughes 6*e1fe3e4aSElliott Hughesclass Visitor(object): 7*e1fe3e4aSElliott Hughes defaultStop = False 8*e1fe3e4aSElliott Hughes 9*e1fe3e4aSElliott Hughes @classmethod 10*e1fe3e4aSElliott Hughes def _register(celf, clazzes_attrs): 11*e1fe3e4aSElliott Hughes assert celf != Visitor, "Subclass Visitor instead." 12*e1fe3e4aSElliott Hughes if "_visitors" not in celf.__dict__: 13*e1fe3e4aSElliott Hughes celf._visitors = {} 14*e1fe3e4aSElliott Hughes 15*e1fe3e4aSElliott Hughes def wrapper(method): 16*e1fe3e4aSElliott Hughes assert method.__name__ == "visit" 17*e1fe3e4aSElliott Hughes for clazzes, attrs in clazzes_attrs: 18*e1fe3e4aSElliott Hughes if type(clazzes) != tuple: 19*e1fe3e4aSElliott Hughes clazzes = (clazzes,) 20*e1fe3e4aSElliott Hughes if type(attrs) == str: 21*e1fe3e4aSElliott Hughes attrs = (attrs,) 22*e1fe3e4aSElliott Hughes for clazz in clazzes: 23*e1fe3e4aSElliott Hughes _visitors = celf._visitors.setdefault(clazz, {}) 24*e1fe3e4aSElliott Hughes for attr in attrs: 25*e1fe3e4aSElliott Hughes assert attr not in _visitors, ( 26*e1fe3e4aSElliott Hughes "Oops, class '%s' has visitor function for '%s' defined already." 27*e1fe3e4aSElliott Hughes % (clazz.__name__, attr) 28*e1fe3e4aSElliott Hughes ) 29*e1fe3e4aSElliott Hughes _visitors[attr] = method 30*e1fe3e4aSElliott Hughes return None 31*e1fe3e4aSElliott Hughes 32*e1fe3e4aSElliott Hughes return wrapper 33*e1fe3e4aSElliott Hughes 34*e1fe3e4aSElliott Hughes @classmethod 35*e1fe3e4aSElliott Hughes def register(celf, clazzes): 36*e1fe3e4aSElliott Hughes if type(clazzes) != tuple: 37*e1fe3e4aSElliott Hughes clazzes = (clazzes,) 38*e1fe3e4aSElliott Hughes return celf._register([(clazzes, (None,))]) 39*e1fe3e4aSElliott Hughes 40*e1fe3e4aSElliott Hughes @classmethod 41*e1fe3e4aSElliott Hughes def register_attr(celf, clazzes, attrs): 42*e1fe3e4aSElliott Hughes clazzes_attrs = [] 43*e1fe3e4aSElliott Hughes if type(clazzes) != tuple: 44*e1fe3e4aSElliott Hughes clazzes = (clazzes,) 45*e1fe3e4aSElliott Hughes if type(attrs) == str: 46*e1fe3e4aSElliott Hughes attrs = (attrs,) 47*e1fe3e4aSElliott Hughes for clazz in clazzes: 48*e1fe3e4aSElliott Hughes clazzes_attrs.append((clazz, attrs)) 49*e1fe3e4aSElliott Hughes return celf._register(clazzes_attrs) 50*e1fe3e4aSElliott Hughes 51*e1fe3e4aSElliott Hughes @classmethod 52*e1fe3e4aSElliott Hughes def register_attrs(celf, clazzes_attrs): 53*e1fe3e4aSElliott Hughes return celf._register(clazzes_attrs) 54*e1fe3e4aSElliott Hughes 55*e1fe3e4aSElliott Hughes @classmethod 56*e1fe3e4aSElliott Hughes def _visitorsFor(celf, thing, _default={}): 57*e1fe3e4aSElliott Hughes typ = type(thing) 58*e1fe3e4aSElliott Hughes 59*e1fe3e4aSElliott Hughes for celf in celf.mro(): 60*e1fe3e4aSElliott Hughes _visitors = getattr(celf, "_visitors", None) 61*e1fe3e4aSElliott Hughes if _visitors is None: 62*e1fe3e4aSElliott Hughes break 63*e1fe3e4aSElliott Hughes 64*e1fe3e4aSElliott Hughes m = celf._visitors.get(typ, None) 65*e1fe3e4aSElliott Hughes if m is not None: 66*e1fe3e4aSElliott Hughes return m 67*e1fe3e4aSElliott Hughes 68*e1fe3e4aSElliott Hughes return _default 69*e1fe3e4aSElliott Hughes 70*e1fe3e4aSElliott Hughes def visitObject(self, obj, *args, **kwargs): 71*e1fe3e4aSElliott Hughes """Called to visit an object. This function loops over all non-private 72*e1fe3e4aSElliott Hughes attributes of the objects and calls any user-registered (via 73*e1fe3e4aSElliott Hughes @register_attr() or @register_attrs()) visit() functions. 74*e1fe3e4aSElliott Hughes 75*e1fe3e4aSElliott Hughes If there is no user-registered visit function, of if there is and it 76*e1fe3e4aSElliott Hughes returns True, or it returns None (or doesn't return anything) and 77*e1fe3e4aSElliott Hughes visitor.defaultStop is False (default), then the visitor will proceed 78*e1fe3e4aSElliott Hughes to call self.visitAttr()""" 79*e1fe3e4aSElliott Hughes 80*e1fe3e4aSElliott Hughes keys = sorted(vars(obj).keys()) 81*e1fe3e4aSElliott Hughes _visitors = self._visitorsFor(obj) 82*e1fe3e4aSElliott Hughes defaultVisitor = _visitors.get("*", None) 83*e1fe3e4aSElliott Hughes for key in keys: 84*e1fe3e4aSElliott Hughes if key[0] == "_": 85*e1fe3e4aSElliott Hughes continue 86*e1fe3e4aSElliott Hughes value = getattr(obj, key) 87*e1fe3e4aSElliott Hughes visitorFunc = _visitors.get(key, defaultVisitor) 88*e1fe3e4aSElliott Hughes if visitorFunc is not None: 89*e1fe3e4aSElliott Hughes ret = visitorFunc(self, obj, key, value, *args, **kwargs) 90*e1fe3e4aSElliott Hughes if ret == False or (ret is None and self.defaultStop): 91*e1fe3e4aSElliott Hughes continue 92*e1fe3e4aSElliott Hughes self.visitAttr(obj, key, value, *args, **kwargs) 93*e1fe3e4aSElliott Hughes 94*e1fe3e4aSElliott Hughes def visitAttr(self, obj, attr, value, *args, **kwargs): 95*e1fe3e4aSElliott Hughes """Called to visit an attribute of an object.""" 96*e1fe3e4aSElliott Hughes self.visit(value, *args, **kwargs) 97*e1fe3e4aSElliott Hughes 98*e1fe3e4aSElliott Hughes def visitList(self, obj, *args, **kwargs): 99*e1fe3e4aSElliott Hughes """Called to visit any value that is a list.""" 100*e1fe3e4aSElliott Hughes for value in obj: 101*e1fe3e4aSElliott Hughes self.visit(value, *args, **kwargs) 102*e1fe3e4aSElliott Hughes 103*e1fe3e4aSElliott Hughes def visitDict(self, obj, *args, **kwargs): 104*e1fe3e4aSElliott Hughes """Called to visit any value that is a dictionary.""" 105*e1fe3e4aSElliott Hughes for value in obj.values(): 106*e1fe3e4aSElliott Hughes self.visit(value, *args, **kwargs) 107*e1fe3e4aSElliott Hughes 108*e1fe3e4aSElliott Hughes def visitLeaf(self, obj, *args, **kwargs): 109*e1fe3e4aSElliott Hughes """Called to visit any value that is not an object, list, 110*e1fe3e4aSElliott Hughes or dictionary.""" 111*e1fe3e4aSElliott Hughes pass 112*e1fe3e4aSElliott Hughes 113*e1fe3e4aSElliott Hughes def visit(self, obj, *args, **kwargs): 114*e1fe3e4aSElliott Hughes """This is the main entry to the visitor. The visitor will visit object 115*e1fe3e4aSElliott Hughes obj. 116*e1fe3e4aSElliott Hughes 117*e1fe3e4aSElliott Hughes The visitor will first determine if there is a registered (via 118*e1fe3e4aSElliott Hughes @register()) visit function for the type of object. If there is, it 119*e1fe3e4aSElliott Hughes will be called, and (visitor, obj, *args, **kwargs) will be passed to 120*e1fe3e4aSElliott Hughes the user visit function. 121*e1fe3e4aSElliott Hughes 122*e1fe3e4aSElliott Hughes If there is no user-registered visit function, of if there is and it 123*e1fe3e4aSElliott Hughes returns True, or it returns None (or doesn't return anything) and 124*e1fe3e4aSElliott Hughes visitor.defaultStop is False (default), then the visitor will proceed 125*e1fe3e4aSElliott Hughes to dispatch to one of self.visitObject(), self.visitList(), 126*e1fe3e4aSElliott Hughes self.visitDict(), or self.visitLeaf() (any of which can be overriden in 127*e1fe3e4aSElliott Hughes a subclass).""" 128*e1fe3e4aSElliott Hughes 129*e1fe3e4aSElliott Hughes visitorFunc = self._visitorsFor(obj).get(None, None) 130*e1fe3e4aSElliott Hughes if visitorFunc is not None: 131*e1fe3e4aSElliott Hughes ret = visitorFunc(self, obj, *args, **kwargs) 132*e1fe3e4aSElliott Hughes if ret == False or (ret is None and self.defaultStop): 133*e1fe3e4aSElliott Hughes return 134*e1fe3e4aSElliott Hughes if hasattr(obj, "__dict__") and not isinstance(obj, enum.Enum): 135*e1fe3e4aSElliott Hughes self.visitObject(obj, *args, **kwargs) 136*e1fe3e4aSElliott Hughes elif isinstance(obj, list): 137*e1fe3e4aSElliott Hughes self.visitList(obj, *args, **kwargs) 138*e1fe3e4aSElliott Hughes elif isinstance(obj, dict): 139*e1fe3e4aSElliott Hughes self.visitDict(obj, *args, **kwargs) 140*e1fe3e4aSElliott Hughes else: 141*e1fe3e4aSElliott Hughes self.visitLeaf(obj, *args, **kwargs) 142