py2dot

1 #!/usr/bin/env python 2 3 """Create a dot file from a Python file. 4 5 Read a Python files and create a dot file of the function used and called. 6 Imported modules are scanned too, up to a user-defined level of recursion. 7 The dot file can be viewed with Graphviz. 8 9 The basic usage is: 10 11 $> ./py2dot -f filename.py | dot -Tpng | display 12 13 For command line options, type: 14 15 $> ./py2dot -h 16 17 Inside a Python shell: 18 19 >>> import py2dot 20 >>> infile = open('test.py', 'r') 21 >>> reclevel = 1 22 >>> data = py2dot.FileInput(infile, maxreclevel=reclevel) 23 >>> print data 24 25 """ 26 27 __author__ = 'Lorenzo Bolla' 28 29 import sys 30 import os.path 31 import parser 32 import symbol 33 import token 34 from pprint import pprint 35 from optparse import OptionParser 36

37 -class HashableName:

38 """Class to redefine the hash function. Used for user-defined data types that must go into sets."""

39 - def __hash__(self):

40 return self.name.__hash__()

41 - def __eq__(self, x):

42 return self.name == x.name

43

44 -class Element:

45 """A generic element of the parsing tree."""

46 - def validate(self, ast):

47 if ast[0] != self.sym_code: 48 raise ValueError('%d != %d' % (ast[0], self.sym_code))

49

50 -class DottedName(Element, HashableName):

51 sym_code = symbol.dotted_name 52

53 - def parse(self, ast):

54 self.validate(ast) 55 name = '' 56 for item in ast[1:]: 57 name += item[1] 58 self.name = name 59 return self

60

61 - def __str__(self):

62 return self.name

63

64 -class DottedAsName(Element, HashableName):

65 sym_code = symbol.dotted_as_name 66

67 - def parse(self, ast):

68 self.validate(ast) 69 dn = DottedName().parse(ast[1]) 70 self.name = dn.name 71 return self

72

73 - def __str__(self):

74 return self.name

75

76 -class DottedAsNames(Element, set):

77 sym_code = symbol.dotted_as_names 78

79 - def parse(self, ast):

80 self.validate(ast) 81 for item in ast[1::2]: 82 das = DottedAsName().parse(item) 83 self.add(das) 84 return self

85

86 -class ImportFrom(Element, set):

87 sym_code = symbol.import_from 88

89 - def parse(self, ast):

90 self.validate(ast) 91 for item in ast[1:]: 92 try: 93 dn = DottedName().parse(item) 94 self.add(dn) 95 except: 96 continue 97 return self

98

99 -class ImportName(Element, set):

100 sym_code = symbol.import_name 101

102 - def parse(self, ast):

103 self.validate(ast) 104 self.update(DottedAsNames().parse(ast[2])) 105 return self

106

107 -class FuncDef(Element, HashableName):

108 sym_code = symbol.funcdef 109

110 - def parse(self, ast):

111 self.validate(ast) 112 self.name = ast[2][1] 113 fc_set = set() 114 recapply(ast[5], FuncCall, fc_set) 115 self.calls = fc_set 116 return self

117

118 -class Atom(Element, HashableName):

119 sym_code = symbol.atom 120

121 - def parse(self, ast):

122 self.validate(ast) 123 if ast[1][0] != token.NAME: 124 raise ValueError() 125 self.name = ast[1][1] 126 return self

127

128 -class FuncCall(Element, HashableName):

129 sym_code = symbol.power 130

131 - def parse(self, ast):

132 global COMPLETE_NAME 133 self.validate(ast) 134 # 'return' statement have no parameters 135 if len(ast) < 3: 136 raise ValueError() 137 self.name = Atom().parse(ast[1]).name 138 for item in ast[2:]: 139 if item[1][0] == token.DOT: 140 # call with full module path 141 if COMPLETE_NAME: 142 self.name = make_name(self.name, item[2][1]) 143 else: 144 # only function name 145 self.name = item[2][1] 146 return self

147

148 -class ClassDef(Element, HashableName):

149 sym_code = symbol.classdef 150

151 - def parse(self, ast):

152 global COMPLETE_NAME 153 self.validate(ast) 154 if ast[2][0] != token.NAME: 155 raise ValueError() 156 self.name = ast[2][1] 157 cl_set = set() 158 for a in ast[4:]: 159 recapply(a, ClassDef, cl_set) 160 self.classes = cl_set 161 fn_set = set() 162 for a in ast[4:]: 163 recapply(a, FuncDef, fn_set) 164 self.funcalls = fn_set 165 # # OKKIO 166 # if COMPLETE_NAME: 167 # for fn in self.funcalls: 168 # fn.name = make_name(self.name, fn.name) 169 # for fc in fn.calls: 170 # fc.name = make_name(self.name, fc.name) 171 return self

172

173 - def relationships(self):

174 dot = '' 175 for fn in self.funcalls: 176 for fc in fn.calls: 177 dot += line('%s -> %s;' % (fn.name, fc.name)) 178 for cl in self.classes: 179 dot += cl.relationships() 180 return dot

181

182 - def definitions(self):

183 dot = '' 184 dot += line('subgraph cluster_%s {' % self.name) 185 dot += line('label = "%s";' % self.name) 186 dot += line('style = dashed;') 187 for fn in self.funcalls: 188 dot += line('%s;' % fn.name) 189 for cl in self.classes: 190 dot += cl.definitions() 191 dot += line('}') 192 return dot

193

194 - def todotlist(self):

195 dot = [] 196 # fun def inside a cluster 197 dot.extend(self.definitions()) 198 # fun calls outside the cluster 199 dot.extend(self.relationships()) 200 return dot

201

202 -class FileInput(Element, HashableName):

203 sym_code = symbol.file_input 204

205 - def __init__(self, filename, reclevel=0, maxreclevel=0):

206 self.reclevel = reclevel 207 self.maxreclevel = maxreclevel 208 if reclevel > maxreclevel: 209 raise ValueError() 210 self.funcalls = set() 211 self.classes = set() 212 self.imports = set() 213 self.file = filename 214 self.name = os.path.basename(self.file.name).split('.')[0] 215 code = self.file.read() 216 217 ast = parser.suite(code) 218 self.parse(ast.tolist())

219

220 - def parse(self, ast):

221 global COMPLETE_NAME 222 global EXCLUDE_CLASSES 223 224 self.validate(ast) 225 226 # classes 227 cl_set = set() 228 if not EXCLUDE_CLASSES: 229 recapply(ast, ClassDef, cl_set) 230 self.classes.update(cl_set) 231 232 # function calls 233 fn_set = set() 234 recapply(ast, FuncDef, fn_set) 235 self.funcalls.update(fn_set) 236 237 # imports 238 import_set = set() 239 recapply(ast, ImportName, import_set) 240 recapply(ast, ImportFrom, import_set) 241 242 path = [os.path.dirname(self.file.name)] + sys.path 243 for im in import_set: 244 for p in path: 245 try: 246 tmpfile = os.path.join(p, im.name + '.py') 247 data = FileInput(open(tmpfile, 'r'), 248 reclevel=self.reclevel+1, 249 maxreclevel=self.maxreclevel) 250 if COMPLETE_NAME: 251 for fn in data.funcalls: 252 fn.name = make_name(im.name, fn.name) 253 for fc in fn.calls: 254 fc.name = make_name(im.name, fc.name) 255 self.imports.add(data) 256 except Exception, e: 257 pass 258 259 return self

260

261 - def definitions(self, subgraph=False):

262 dot = '' 263 dot += line('/* definitions */') 264 if subgraph: 265 dot += line('subgraph cluster_%s {' % self.name) 266 dot += line('style = filled; fillcolor = lightgrey;') 267 dot += line('label = "%s";' % self.name) 268 dot += line('splines=true;') 269 dot += line('size="7,7";') 270 for cl in self.classes: 271 dot += cl.definitions() 272 for fn in self.funcalls: 273 dot += line('%s;' % fn.name) 274 for imp in self.imports: 275 dot += imp.definitions(subgraph=True) 276 if subgraph: 277 dot += line('}') 278 return dot

279

280 - def relationships(self):

281 dot = '' 282 dot += line('/* relationships */') 283 for fn in self.funcalls: 284 for fc in fn.calls: 285 dot += line('%s -> %s;' % (fn.name, fc.name)) 286 for cl in self.classes: 287 dot += cl.relationships() 288 for imp in self.imports: 289 dot += imp.relationships() 290 return dot

291

292 - def open(self):

293 return line('digraph %s {' % self.name)

294

295 - def close(self):

296 return line('}')

297

298 - def todot(self):

299 dot = '' 300 dot += self.open() 301 dot += self.definitions() 302 dot += self.relationships() 303 dot += self.close() 304 return dot

305

306 - def __str__(self):

307 return self.todot()

308

309 -def make_name(x, y, sep='_'):

310 return x + sep + y

311

312 -def line(x):

313 return x + '\n'

314

315 -def isiterable(x):

316 return isinstance(x, (tuple, list, set))

317

318 -def recapply(ast, cls, data):

319 """Iteratively apply a class constructor to the parser list.""" 320 try: 321 tmp = cls().parse(ast) 322 del ast[0] 323 if isiterable(tmp): 324 data.update(tmp) 325 else: 326 data.add(tmp) 327 328 except: 329 pass 330 331 finally: 332 for item in ast[1:]: 333 if isiterable(item): 334 recapply(item, cls, data)

335 336 if __name__ == '__main__': 337 global COMPLETE_NAME 338 global EXCLUDE_CLASSES 339 340 oparser = OptionParser() 341 oparser.add_option('-f', '--file', dest='infile', default=sys.stdin, type='string', 342 help='Input file [stdin by default]') 343 oparser.add_option('-o', '--output', dest='outfile', default=sys.stdout, type='string', 344 help='Output file [stdout by default]') 345 oparser.add_option('-i', '--incomplete_name', dest='incomplete_name', default=False, action='store_true', 346 help='Do not prepend module names to imported functions') 347 oparser.add_option('-x', '--exclude_classes', dest='exclude_classes', default=False, action='store_true', 348 help='Do not create a subgaph for classes') 349 oparser.add_option('-r', '--recursion_level', dest='reclevel', default=0, type='int', 350 help='Maximum level of recursion in scanning imported modules') 351 352 (options, args) = oparser.parse_args() 353 354 infile = options.infile 355 if isinstance(infile, str): 356 infile = open(infile, 'r') 357 358 outfile = options.outfile 359 if isinstance(outfile, str): 360 outfile = open(outfile, 'w') 361 362 COMPLETE_NAME = not options.incomplete_name 363 EXCLUDE_CLASSES = options.exclude_classes 364 reclevel = options.reclevel 365 366 data = FileInput(infile, maxreclevel=reclevel) 367 368 outfile.write(str(data)) 369

Source Code for Module py2dot