| """Generic visitor pattern implementation for Python objects.""" |
| |
| import enum |
| |
| |
| class Visitor(object): |
| defaultStop = False |
| |
| @classmethod |
| def _register(celf, clazzes_attrs): |
| assert celf != Visitor, "Subclass Visitor instead." |
| if "_visitors" not in celf.__dict__: |
| celf._visitors = {} |
| |
| def wrapper(method): |
| assert method.__name__ == "visit" |
| for clazzes, attrs in clazzes_attrs: |
| if type(clazzes) != tuple: |
| clazzes = (clazzes,) |
| if type(attrs) == str: |
| attrs = (attrs,) |
| for clazz in clazzes: |
| _visitors = celf._visitors.setdefault(clazz, {}) |
| for attr in attrs: |
| assert attr not in _visitors, ( |
| "Oops, class '%s' has visitor function for '%s' defined already." |
| % (clazz.__name__, attr) |
| ) |
| _visitors[attr] = method |
| return None |
| |
| return wrapper |
| |
| @classmethod |
| def register(celf, clazzes): |
| if type(clazzes) != tuple: |
| clazzes = (clazzes,) |
| return celf._register([(clazzes, (None,))]) |
| |
| @classmethod |
| def register_attr(celf, clazzes, attrs): |
| clazzes_attrs = [] |
| if type(clazzes) != tuple: |
| clazzes = (clazzes,) |
| if type(attrs) == str: |
| attrs = (attrs,) |
| for clazz in clazzes: |
| clazzes_attrs.append((clazz, attrs)) |
| return celf._register(clazzes_attrs) |
| |
| @classmethod |
| def register_attrs(celf, clazzes_attrs): |
| return celf._register(clazzes_attrs) |
| |
| @classmethod |
| def _visitorsFor(celf, thing, _default={}): |
| typ = type(thing) |
| |
| for celf in celf.mro(): |
| _visitors = getattr(celf, "_visitors", None) |
| if _visitors is None: |
| break |
| |
| m = celf._visitors.get(typ, None) |
| if m is not None: |
| return m |
| |
| return _default |
| |
| def visitObject(self, obj, *args, **kwargs): |
| """Called to visit an object. This function loops over all non-private |
| attributes of the objects and calls any user-registered (via |
| @register_attr() or @register_attrs()) visit() functions. |
| |
| If there is no user-registered visit function, of if there is and it |
| returns True, or it returns None (or doesn't return anything) and |
| visitor.defaultStop is False (default), then the visitor will proceed |
| to call self.visitAttr()""" |
| |
| keys = sorted(vars(obj).keys()) |
| _visitors = self._visitorsFor(obj) |
| defaultVisitor = _visitors.get("*", None) |
| for key in keys: |
| if key[0] == "_": |
| continue |
| value = getattr(obj, key) |
| visitorFunc = _visitors.get(key, defaultVisitor) |
| if visitorFunc is not None: |
| ret = visitorFunc(self, obj, key, value, *args, **kwargs) |
| if ret == False or (ret is None and self.defaultStop): |
| continue |
| self.visitAttr(obj, key, value, *args, **kwargs) |
| |
| def visitAttr(self, obj, attr, value, *args, **kwargs): |
| """Called to visit an attribute of an object.""" |
| self.visit(value, *args, **kwargs) |
| |
| def visitList(self, obj, *args, **kwargs): |
| """Called to visit any value that is a list.""" |
| for value in obj: |
| self.visit(value, *args, **kwargs) |
| |
| def visitDict(self, obj, *args, **kwargs): |
| """Called to visit any value that is a dictionary.""" |
| for value in obj.values(): |
| self.visit(value, *args, **kwargs) |
| |
| def visitLeaf(self, obj, *args, **kwargs): |
| """Called to visit any value that is not an object, list, |
| or dictionary.""" |
| pass |
| |
| def visit(self, obj, *args, **kwargs): |
| """This is the main entry to the visitor. The visitor will visit object |
| obj. |
| |
| The visitor will first determine if there is a registered (via |
| @register()) visit function for the type of object. If there is, it |
| will be called, and (visitor, obj, *args, **kwargs) will be passed to |
| the user visit function. |
| |
| If there is no user-registered visit function, of if there is and it |
| returns True, or it returns None (or doesn't return anything) and |
| visitor.defaultStop is False (default), then the visitor will proceed |
| to dispatch to one of self.visitObject(), self.visitList(), |
| self.visitDict(), or self.visitLeaf() (any of which can be overriden in |
| a subclass).""" |
| |
| visitorFunc = self._visitorsFor(obj).get(None, None) |
| if visitorFunc is not None: |
| ret = visitorFunc(self, obj, *args, **kwargs) |
| if ret == False or (ret is None and self.defaultStop): |
| return |
| if hasattr(obj, "__dict__") and not isinstance(obj, enum.Enum): |
| self.visitObject(obj, *args, **kwargs) |
| elif isinstance(obj, list): |
| self.visitList(obj, *args, **kwargs) |
| elif isinstance(obj, dict): |
| self.visitDict(obj, *args, **kwargs) |
| else: |
| self.visitLeaf(obj, *args, **kwargs) |
