| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158 |
- """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}
- }
- @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
- 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)
|