|
1 | 1 | import ast |
| 2 | +import random |
2 | 3 | from collections import namedtuple |
3 | 4 |
|
4 | 5 | from .node_types import ( |
@@ -101,16 +102,24 @@ def get_first_node( |
101 | 102 | node, |
102 | 103 | node_not_to_step_past |
103 | 104 | ): |
| 105 | + """ |
| 106 | + This is a super hacky way of getting the first node after a statement. |
| 107 | + We do this because we visit a statement and keep on visiting and get something in return that is rarely the first node. |
| 108 | + So we loop and loop backwards until we hit the statement or there is nothing to step back to. |
| 109 | + """ |
104 | 110 | ingoing = None |
| 111 | + i = 0 |
105 | 112 | current_node = node |
106 | 113 | while current_node.ingoing: |
| 114 | + # This is used because there may be multiple ingoing and loop will cause an infinite loop if we did [0] |
| 115 | + i = random.randrange(len(current_node.ingoing)) |
107 | 116 | # e.g. We don't want to step past the Except of an Except basic block |
108 | | - if current_node.ingoing[0] == node_not_to_step_past: |
| 117 | + if current_node.ingoing[i] == node_not_to_step_past: |
109 | 118 | break |
110 | 119 | ingoing = current_node.ingoing |
111 | | - current_node = current_node.ingoing[0] |
| 120 | + current_node = current_node.ingoing[i] |
112 | 121 | if ingoing: |
113 | | - return ingoing[0] |
| 122 | + return ingoing[i] |
114 | 123 | return current_node |
115 | 124 |
|
116 | 125 |
|
|
0 commit comments