Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2@file 

3@brief One class which visits a syntax tree. 

4""" 

5 

6import ast 

7 

8 

9class CodeNodeVisitor(ast.NodeVisitor): 

10 

11 """ 

12 Defines a visitor which walks though the syntax tree of the code. 

13 

14 .. exref:: 

15 :title: Get the tree of a simple function 

16 

17 The following code uses Python syntax but follows a SQL logic. 

18 

19 .. runpython:: 

20 :showcode: 

21 :process: 

22 

23 import ast 

24 import inspect 

25 import textwrap 

26 from pysqllike.translation.node_visitor_translator import CodeNodeVisitor 

27 

28 def myjob(input): 

29 iter = input.select (input.age, input.nom, age2 = input.age2*input.age2) 

30 wher = iter.where( (iter.age > 60).Or(iter.age < 25)) 

31 return wher 

32 

33 code = textwrap.dedent(inspect.getsource(myjob)) 

34 node = ast.parse(code) 

35 v = CodeNodeVisitor() 

36 v.visit(node) 

37 for r in v.Rows : 

38 print("{0}{1}: {2}".format(" " * r["indent"], r["type"], r["str"])) 

39 """ 

40 

41 def __init__(self): 

42 """ 

43 constructor 

44 """ 

45 ast.NodeVisitor.__init__(self) 

46 self._rows = [] 

47 self._indent = 0 

48 self._stack = [] 

49 

50 def push(self, row): 

51 """ 

52 Pushes an element into a list. 

53 """ 

54 self._rows.append(row) 

55 

56 def generic_visit(self, node): 

57 """ 

58 Overrides ``generic_visit`` to check it is not used. 

59 """ 

60 raise AttributeError("generic_visit_args should be used.") 

61 

62 def generic_visit_args(self, node, row): 

63 """ 

64 Overrides ``generic_visit`` to keep track of the indentation 

65 and the node parent. The function will add field 

66 ``row["children"] = visited`` nodes from here. 

67 

68 @param node node which needs to be visited 

69 @param row row (a dictionary) 

70 @return See ``ast.NodeVisitor.generic_visit`` 

71 """ 

72 self._indent += 1 

73 last = len(self._rows) 

74 res = ast.NodeVisitor.generic_visit(self, node) 

75 row["children"] = [ 

76 _ for _ in self._rows[ 

77 last:] if _["indent"] == self._indent] 

78 self._indent -= 1 

79 return res 

80 

81 def visit(self, node): 

82 """ 

83 Visits a node, a method must exist for every object class. 

84 """ 

85 method = 'visit_' + node.__class__.__name__ 

86 visitor = getattr(self, method, None) 

87 if visitor is None: 

88 raise TypeError("unable to find a method: " + method) 

89 res = visitor(node) 

90 # print(method, CodeNodeVisitor.print_node(node)) 

91 return res 

92 

93 @staticmethod 

94 def print_node(node): 

95 """ 

96 Debugging purpose. 

97 """ 

98 r = [] 

99 for att in ["s", "name", "str", "id", "body", "n", 

100 "arg", "targets", "attr", "returns", "ctx"]: 

101 if att in node.__dict__: 

102 r.append("{0}={1}".format(att, str(node.__dict__[att]))) 

103 return " ".join(r) 

104 

105 def print_tree(self): 

106 """ 

107 Displays the tree of instructions. 

108 

109 @return string 

110 """ 

111 rows = [] 

112 for r in self.Rows: 

113 rows.append( 

114 ("{0}{1}: {2}".format( 

115 " " * 

116 r["indent"], 

117 r["type"], 

118 r["str"]))) 

119 return "\n".join(rows) 

120 

121 @property 

122 def Rows(self): 

123 """ 

124 returns a list of dictionaries with all the elements of the code 

125 """ 

126 return [_ for _ in self._rows if not _.get("remove", False)] 

127 

128 def visit_Str(self, node): 

129 cont = { 

130 "indent": self._indent, 

131 "type": "Str", 

132 "str": node.s, 

133 "node": node, 

134 "value": node.s} 

135 self.push(cont) 

136 return self.generic_visit_args(node, cont) 

137 

138 def visit_Name(self, node): 

139 cont = { 

140 "indent": self._indent, 

141 "type": "Name", 

142 "str": node.id, 

143 "node": node, 

144 "id": node.id, 

145 "ctx": node.ctx} 

146 self.push(cont) 

147 return self.generic_visit_args(node, cont) 

148 

149 def visit_Module(self, node): 

150 cont = { 

151 "indent": self._indent, 

152 "type": "Module", 

153 "str": "", 

154 "body": node.body, 

155 "node": node} 

156 self.push(cont) 

157 return self.generic_visit_args(node, cont) 

158 

159 def visit_FunctionDef(self, node): 

160 cont = {"indent": self._indent, "type": "FunctionDef", "str": node.name, "name": node.name, "body": node.body, 

161 "node": node, "returns": node.returns} 

162 self.push(cont) 

163 return self.generic_visit_args(node, cont) 

164 

165 def visit_arguments(self, node): 

166 cont = {"indent": self._indent, "type": "arguments", "str": "", 

167 "node": node, "args": node.args} 

168 self.push(cont) 

169 return self.generic_visit_args(node, cont) 

170 

171 def visit_arg(self, node): 

172 cont = {"indent": self._indent, "type": "arg", "str": node.arg, 

173 "node": node, 

174 "arg": node.arg, "annotation": node.annotation} 

175 self.push(cont) 

176 return self.generic_visit_args(node, cont) 

177 

178 def visit_Assign(self, node): 

179 cont = {"indent": self._indent, "type": "Assign", "str": "", "node": node, 

180 "targets": node.targets, "value": node.value} 

181 self.push(cont) 

182 return self.generic_visit_args(node, cont) 

183 

184 def visit_Store(self, node): 

185 #cont = { "indent":self._indent, "type": "Store", "str": "" } 

186 # self.push(cont) 

187 cont = {} 

188 return self.generic_visit_args(node, cont) 

189 

190 def visit_Call(self, node): 

191 if "attr" in node.func.__dict__: 

192 cont = {"indent": self._indent, "type": "Call", "str": node.func.attr, 

193 "node": node, "func": node.func} 

194 else: 

195 cont = {"indent": self._indent, "type": "Call", "str": node.func.id, 

196 "node": node, "func": node.func} 

197 self.push(cont) 

198 return self.generic_visit_args(node, cont) 

199 

200 def visit_Attribute(self, node): 

201 cont = {"indent": self._indent, "type": "Attribute", "str": node.attr, 

202 "node": node, "value": node.value, "ctx": node.ctx, "attr": node.attr} 

203 self.push(cont) 

204 # last = len(self._rows) 

205 res = self.generic_visit_args(node, cont) 

206 

207 if len(cont["children"]) > 0: 

208 fir = cont["children"][0] 

209 if fir["type"] == "Name": 

210 parent = fir["node"].id 

211 cont["str"] = "{0}.{1}".format(parent, cont["str"]) 

212 cont["children"][0]["remove"] = True 

213 return res 

214 

215 def visit_Load(self, node): 

216 #cont = { "indent":self._indent, "type": "Load", "str": "" } 

217 # self.push(cont) 

218 cont = {} 

219 return self.generic_visit_args(node, cont) 

220 

221 def visit_keyword(self, node): 

222 cont = {"indent": self._indent, "type": "keyword", "str": "{0}".format(node.arg), 

223 "node": node, "arg": node.arg, "value": node.value} 

224 self.push(cont) 

225 return self.generic_visit_args(node, cont) 

226 

227 def visit_BinOp(self, node): 

228 cont = { 

229 "indent": self._indent, 

230 "type": "BinOp", 

231 "str": "", 

232 "node": node} 

233 self.push(cont) 

234 return self.generic_visit_args(node, cont) 

235 

236 def visit_Mult(self, node): 

237 cont = { 

238 "indent": self._indent, 

239 "type": "Mult", 

240 "str": "", 

241 "node": node} 

242 self.push(cont) 

243 return self.generic_visit_args(node, cont) 

244 

245 def visit_Compare(self, node): 

246 cont = { 

247 "indent": self._indent, 

248 "type": "Compare", 

249 "str": "", 

250 "node": node} 

251 self.push(cont) 

252 return self.generic_visit_args(node, cont) 

253 

254 def visit_Gt(self, node): 

255 cont = {"indent": self._indent, "type": "Gt", "str": "", "node": node} 

256 self.push(cont) 

257 return self.generic_visit_args(node, cont) 

258 

259 def visit_Lt(self, node): 

260 cont = {"indent": self._indent, "type": "Lt", "str": "", "node": node} 

261 self.push(cont) 

262 return self.generic_visit_args(node, cont) 

263 

264 def visit_Num(self, node): 

265 cont = { 

266 "indent": self._indent, 

267 "type": "Num", 

268 "node": node, 

269 "str": "{0}".format( 

270 node.n), 

271 'n': node.n} 

272 self.push(cont) 

273 return self.generic_visit_args(node, cont) 

274 

275 def visit_Return(self, node): 

276 cont = {"indent": self._indent, "type": "Return", "node": node, "str": "", 

277 'value': node.value} 

278 self.push(cont) 

279 return self.generic_visit_args(node, cont) 

280 

281 def visit_(self, node): 

282 help(node) 

283 assert False