Spaces:
Running
Running
| """Generic visitor pattern implementation for Python objects.""" | |
| import enum | |
| import weakref | |
| class Visitor(object): | |
| defaultStop = False | |
| _visitors = { | |
| # By default we skip visiting weak references to avoid recursion | |
| # issues. Users can override this by registering a visit | |
| # function for weakref.ProxyType. | |
| weakref.ProxyType: {None: lambda self, obj, *args, **kwargs: False} | |
| } | |
| 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 | |
| def register(celf, clazzes): | |
| if type(clazzes) != tuple: | |
| clazzes = (clazzes,) | |
| return celf._register([(clazzes, (None,))]) | |
| 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) | |
| def register_attrs(celf, clazzes_attrs): | |
| return celf._register(clazzes_attrs) | |
| def _visitorsFor(celf, thing, _default={}): | |
| typ = type(thing) | |
| for celf in celf.mro(): | |
| _visitors = getattr(celf, "_visitors", None) | |
| if _visitors is None: | |
| break | |
| for base in typ.mro(): | |
| m = celf._visitors.get(base, 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. | |
| The visitor will proceed to call ``self.visitAttr()``, unless there is a | |
| user-registered visit function and: | |
| * It returns ``False``; or | |
| * It returns ``None`` (or doesn't return anything) and | |
| ``visitor.defaultStop`` is ``True`` (non-default). | |
| """ | |
| 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. | |
| The visitor will not recurse if there is a user-registered visit | |
| function and: | |
| * It returns ``False``; or | |
| * It returns ``None`` (or doesn't return anything) and | |
| ``visitor.defaultStop`` is ``True`` (non-default) | |
| Otherwise, 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) | |