|
5 | 5 | from . import ast
|
6 | 6 | from .visitor_meta import QUERY_DOCUMENT_KEYS, VisitorMeta
|
7 | 7 |
|
| 8 | + |
| 9 | +class Falsey(object): |
| 10 | + def __nonzero__(self): |
| 11 | + return False |
| 12 | + |
| 13 | + def __bool__(self): |
| 14 | + return False |
| 15 | + |
| 16 | + |
8 | 17 | BREAK = object()
|
9 |
| -REMOVE = object() |
| 18 | +REMOVE = Falsey() |
10 | 19 |
|
11 | 20 |
|
12 | 21 | class Stack(object):
|
@@ -90,7 +99,7 @@ def visit(root, visitor, key_map=None):
|
90 | 99 | key = None
|
91 | 100 | node = new_root
|
92 | 101 |
|
93 |
| - if node is None: |
| 102 | + if node is REMOVE or node is None: |
94 | 103 | continue
|
95 | 104 |
|
96 | 105 | if parent:
|
@@ -162,3 +171,55 @@ def leave(self, node, key, parent, path, ancestors):
|
162 | 171 | method = self._get_leave_handler(type(node))
|
163 | 172 | if method:
|
164 | 173 | return method(self, node, key, parent, path, ancestors)
|
| 174 | + |
| 175 | + |
| 176 | +class ParallelVisitor(Visitor): |
| 177 | + __slots__ = 'skipping', 'visitors' |
| 178 | + |
| 179 | + def __init__(self, visitors): |
| 180 | + self.visitors = visitors |
| 181 | + self.skipping = [None] * len(visitors) |
| 182 | + |
| 183 | + def enter(self, node, key, parent, path, ancestors): |
| 184 | + for i, visitor in enumerate(self.visitors): |
| 185 | + if not self.skipping[i]: |
| 186 | + result = visitor.enter(node, key, parent, path, ancestors) |
| 187 | + if result is False: |
| 188 | + self.skipping[i] = node |
| 189 | + elif result is BREAK: |
| 190 | + self.skipping[i] = BREAK |
| 191 | + elif result is not None: |
| 192 | + return result |
| 193 | + |
| 194 | + def leave(self, node, key, parent, path, ancestors): |
| 195 | + for i, visitor in enumerate(self.visitors): |
| 196 | + if not self.skipping[i]: |
| 197 | + result = visitor.leave(node, key, parent, path, ancestors) |
| 198 | + if result is BREAK: |
| 199 | + self.skipping[i] = BREAK |
| 200 | + elif result is not None and result is not False: |
| 201 | + return result |
| 202 | + elif self.skipping[i] == node: |
| 203 | + self.skipping[i] = REMOVE |
| 204 | + |
| 205 | + |
| 206 | +class TypeInfoVisitor(Visitor): |
| 207 | + __slots__ = 'visitor', 'type_info' |
| 208 | + |
| 209 | + def __init__(self, type_info, visitor): |
| 210 | + self.type_info = type_info |
| 211 | + self.visitor = visitor |
| 212 | + |
| 213 | + def enter(self, node, key, parent, path, ancestors): |
| 214 | + self.type_info.enter(node) |
| 215 | + result = self.visitor.enter(node, key, parent, path, ancestors) |
| 216 | + if result is not None: |
| 217 | + self.type_info.leave(node) |
| 218 | + if isinstance(result, ast.Node): |
| 219 | + self.type_info.enter(result) |
| 220 | + return result |
| 221 | + |
| 222 | + def leave(self, node, key, parent, path, ancestors): |
| 223 | + result = self.visitor.leave(node, key, parent, path, ancestors) |
| 224 | + self.type_info.leave(node) |
| 225 | + return result |
0 commit comments