/ astmonkey / visitors.py
visitors.py
   1  import ast
   2  import re
   3  from contextlib import contextmanager
   4  
   5  import pydot
   6  
   7  from astmonkey import utils
   8  from astmonkey.transformers import ParentChildNodeTransformer
   9  from astmonkey.utils import CommaWriter, check_version
  10  
  11  
  12  class GraphNodeVisitor(ast.NodeVisitor):
  13      def __init__(self):
  14          self.graph = pydot.Dot(graph_type='graph', **self._dot_graph_kwargs())
  15  
  16      def visit(self, node):
  17          if len(node.parents) <= 1:
  18              self.graph.add_node(self._dot_node(node))
  19          if len(node.parents) == 1:
  20              self.graph.add_edge(self._dot_edge(node))
  21          super(GraphNodeVisitor, self).visit(node)
  22  
  23      def _dot_graph_kwargs(self):
  24          return {}
  25  
  26      def _dot_node(self, node):
  27          return pydot.Node(str(node), label=self._dot_node_label(node), **self._dot_node_kwargs(node))
  28  
  29      def _dot_node_label(self, node):
  30          fields_labels = []
  31          for field, value in ast.iter_fields(node):
  32              if not isinstance(value, list):
  33                  value_label = self._dot_node_value_label(value)
  34                  if value_label:
  35                      fields_labels.append('{0}={1}'.format(field, value_label))
  36          return 'ast.{0}({1})'.format(node.__class__.__name__, ', '.join(fields_labels))
  37  
  38      def _dot_node_value_label(self, value):
  39          if not isinstance(value, ast.AST):
  40              return repr(value)
  41          elif len(value.parents) > 1:
  42              return self._dot_node_label(value)
  43          return None
  44  
  45      def _dot_node_kwargs(self, node):
  46          return {
  47              'shape': 'box',
  48              'fontname': 'Curier'
  49          }
  50  
  51      def _dot_edge(self, node):
  52          return pydot.Edge(str(node.parent), str(node), label=self._dot_edge_label(node), **self._dot_edge_kwargs(node))
  53  
  54      def _dot_edge_label(self, node):
  55          label = node.parent_field
  56          if not node.parent_field_index is None:
  57              label += '[{0}]'.format(node.parent_field_index)
  58          return label
  59  
  60      def _dot_edge_kwargs(self, node):
  61          return {
  62              'fontname': 'Curier'
  63          }
  64  
  65  
  66  """
  67  Source generator node visitor from Python AST was originaly written by Armin Ronacher (2008), license BSD.
  68  """
  69  
  70  BOOLOP_SYMBOLS = {
  71      ast.And: 'and',
  72      ast.Or: 'or'
  73  }
  74  
  75  BINOP_SYMBOLS = {
  76      ast.Add: '+',
  77      ast.Sub: '-',
  78      ast.Mult: '*',
  79      ast.Div: '/',
  80      ast.FloorDiv: '//',
  81      ast.Mod: '%',
  82      ast.LShift: '<<',
  83      ast.RShift: '>>',
  84      ast.BitOr: '|',
  85      ast.BitAnd: '&',
  86      ast.BitXor: '^',
  87      ast.Pow: '**'
  88  }
  89  
  90  if check_version(from_inclusive=(3, 5)):
  91      BINOP_SYMBOLS[ast.MatMult] = '@'
  92  
  93  CMPOP_SYMBOLS = {
  94      ast.Eq: '==',
  95      ast.Gt: '>',
  96      ast.GtE: '>=',
  97      ast.In: 'in',
  98      ast.Is: 'is',
  99      ast.IsNot: 'is not',
 100      ast.Lt: '<',
 101      ast.LtE: '<=',
 102      ast.NotEq: '!=',
 103      ast.NotIn: 'not in'
 104  }
 105  
 106  UNARYOP_SYMBOLS = {
 107      ast.Invert: '~',
 108      ast.Not: 'not',
 109      ast.UAdd: '+',
 110      ast.USub: '-'
 111  }
 112  
 113  ALL_SYMBOLS = {}
 114  ALL_SYMBOLS.update(BOOLOP_SYMBOLS)
 115  ALL_SYMBOLS.update(BINOP_SYMBOLS)
 116  ALL_SYMBOLS.update(CMPOP_SYMBOLS)
 117  ALL_SYMBOLS.update(UNARYOP_SYMBOLS)
 118  
 119  
 120  def to_source(node, indent_with=' ' * 4):
 121      """This function can convert a node tree back into python sourcecode.
 122      This is useful for debugging purposes, especially if you're dealing with
 123      custom asts not generated by python itself.
 124  
 125      It could be that the sourcecode is evaluable when the AST itself is not
 126      compilable / evaluable.  The reason for this is that the AST contains some
 127      more data than regular sourcecode does, which is dropped during
 128      conversion.
 129  
 130      Each level of indentation is replaced with `indent_with`.  Per default this
 131      parameter is equal to four spaces as suggested by PEP 8, but it might be
 132      adjusted to match the application's styleguide.
 133      """
 134      ParentChildNodeTransformer().visit(node)
 135      FixLinenoNodeVisitor().visit(node)
 136      generator = SourceGeneratorNodeVisitor(indent_with)
 137      generator.visit(node)
 138  
 139      return ''.join(generator.result)
 140  
 141  
 142  class FixLinenoNodeVisitor(ast.NodeVisitor):
 143      """A helper node visitor for the SourceGeneratorNodeVisitor.
 144  
 145      Attempts to correct implausible line numbers. An example would be:
 146  
 147      1: while a:
 148      2:   pass
 149      3: for a:
 150      2:   pass
 151  
 152      This would be corrected to:
 153  
 154      1: while a:
 155      2:   pass
 156      3: for a:
 157      4:   pass
 158      """
 159  
 160      def __init__(self):
 161          self.min_lineno = 0
 162  
 163      def generic_visit(self, node):
 164          if hasattr(node, 'lineno'):
 165              self._fix_lineno(node)
 166          if hasattr(node, 'body') and isinstance(node.body, list):
 167              self._process_body(node)
 168  
 169      def _fix_lineno(self, node):
 170          if node.lineno < self.min_lineno:
 171              node.lineno = self.min_lineno
 172          else:
 173              self.min_lineno = node.lineno
 174  
 175      def _process_body(self, node):
 176          for body_node in node.body:
 177              self.min_lineno += 1
 178              self.visit(body_node)
 179  
 180  
 181  class BaseSourceGeneratorNodeVisitor(ast.NodeVisitor):
 182      """This visitor is able to transform a well formed syntax tree into python
 183      sourcecode.  For more details have a look at the docstring of the
 184      `node_to_source` function.
 185      """
 186  
 187      def __init__(self, indent_with):
 188          self.result = []
 189          self.indent_with = indent_with
 190          self.indentation = 0
 191  
 192      @classmethod
 193      def _is_node_args_valid(cls, node, arg_name):
 194          return hasattr(node, arg_name) and getattr(node, arg_name) is not None
 195  
 196      def _get_current_line_no(self):
 197          lines = len("".join(self.result).split('\n')) if self.result else 0
 198          return lines
 199  
 200      @classmethod
 201      def _get_actual_lineno(cls, node):
 202          if isinstance(node, (ast.Expr, ast.Str)) and node.col_offset == -1:
 203              str_content = cls._get_string_content(node)
 204              node_lineno = node.lineno - str_content.count('\n')
 205          else:
 206              node_lineno = node.lineno
 207          return node_lineno
 208  
 209      @staticmethod
 210      def _get_string_content(node):
 211          # node is a multi line string and the line number is actually the last line
 212          if isinstance(node, ast.Expr):
 213              str_content = node.value.s
 214          else:
 215              str_content = node.s
 216          if type(str_content) == bytes:
 217              str_content = str_content.decode("utf-8")
 218          return str_content
 219  
 220      def _newline_needed(self, node):
 221          lines = self._get_current_line_no()
 222          node_lineno = self._get_actual_lineno(node)
 223          line_diff = node_lineno - lines
 224          return line_diff > 0
 225  
 226      @contextmanager
 227      def indent(self, count=1):
 228          self.indentation += count
 229          yield
 230          self.indentation -= count
 231  
 232      @contextmanager
 233      def inside(self, pre, post, cond=True):
 234          if cond:
 235              self.write(pre)
 236          yield
 237          if cond:
 238              self.write(post)
 239  
 240      def write(self, x):
 241          self.result.append(x)
 242  
 243      def correct_line_number(self, node, within_statement=True, use_line_continuation=True):
 244          if not node or not self._is_node_args_valid(node, 'lineno'):
 245              return
 246          if within_statement:
 247              indent = 1
 248          else:
 249              indent = 0
 250          with self.indent(indent):
 251              self.add_missing_lines(node, within_statement, use_line_continuation)
 252  
 253      def add_missing_lines(self, node, within_statement, use_line_continuation):
 254          while self._newline_needed(node):
 255              self.add_line(within_statement, use_line_continuation)
 256  
 257      def add_line(self, within_statement, use_line_continuation):
 258          if within_statement and use_line_continuation:
 259              self.result.append('\\')
 260          self.write_newline()
 261  
 262      def write_newline(self):
 263          if self.result:
 264              self.result.append('\n')
 265          self.result.append(self.indent_with * self.indentation)
 266  
 267      def body(self, statements, indent=1):
 268          if statements:
 269              with self.indent(indent):
 270                  for stmt in statements:
 271                      self.correct_line_number(stmt, within_statement=False)
 272                      self.visit(stmt)
 273  
 274      def body_or_else(self, node):
 275          self.body(node.body)
 276          if node.orelse:
 277              self.or_else(node)
 278  
 279      def keyword_and_body(self, keyword, body):
 280          if self._newline_needed(body[0]):
 281              self.write_newline()
 282          self.write(keyword)
 283          self.body(body)
 284  
 285      def or_else(self, node):
 286          self.keyword_and_body('else:', node.orelse)
 287  
 288      def docstring(self, node):
 289          s = repr(node.s)
 290          s = re.sub(r'(?<!\\)\\n', '\n', s)
 291          s = re.sub(r'(?<!\\)\\t', '\t', s)
 292          self.write('%s%s%s' % (s[0] * 2, s, s[0] * 2))
 293  
 294      def signature(self, node, add_space=False):
 295          write_comma = CommaWriter(self.write, add_space_at_beginning=add_space)
 296          padding = [None] * (len(node.args) - len(node.defaults))
 297  
 298          for arg, default in zip(node.args, padding + node.defaults):
 299              self.signature_arg(arg, default, write_comma)
 300  
 301          self.signature_spec_arg(node, 'vararg', write_comma, prefix='*')
 302          self.signature_kwonlyargs(node, write_comma)
 303          self.signature_spec_arg(node, 'kwarg', write_comma, prefix='**')
 304  
 305      def signature_arg(self, arg, default, write_comma, prefix=''):
 306          write_comma()
 307          self.write(prefix)
 308          self.visit(arg)
 309  
 310          if self._is_node_args_valid(arg, 'annotation'):
 311              self.write(': ')
 312              self.visit(arg.annotation)
 313              if default is not None:
 314                  self.write(' = ')
 315                  self.visit(default)
 316          elif default is not None:
 317              self.write('=')
 318              self.visit(default)
 319  
 320      def signature_kwonlyargs(self, node, write_comma):
 321          if not self._is_node_args_valid(node, 'kwonlyargs') or len(node.kwonlyargs) == 0:
 322              return
 323  
 324          if not node.vararg:
 325              write_comma()
 326              self.write('*')
 327  
 328          for arg, default in zip(node.kwonlyargs, node.kw_defaults):
 329              self.signature_arg(arg, default, write_comma)
 330  
 331      def signature_spec_arg(self, node, var, write_comma, prefix):
 332          arg = getattr(node, var)
 333          if arg:
 334              if hasattr(node, var + 'annotation'):
 335                  arg = ast.arg(arg, getattr(node, var + 'annotation'))
 336              self.signature_arg(arg, None, write_comma, prefix)
 337  
 338      def decorators(self, node):
 339          if node.decorator_list:
 340              for decorator in node.decorator_list:
 341                  self.write('@')
 342                  self.visit(decorator)
 343              self.write_newline()
 344  
 345      def visit(self, node):
 346          self.correct_line_number(node)
 347          return super(BaseSourceGeneratorNodeVisitor, self).visit(node)
 348  
 349      # Statements
 350  
 351      def visit_Module(self, node):
 352          self.body(node.body, indent=0)
 353  
 354      def visit_Assign(self, node):
 355  
 356          for idx, target in enumerate(node.targets):
 357              if idx:
 358                  self.write(' = ')
 359              self.visit(target)
 360          self.write(' = ')
 361          self.visit(node.value)
 362  
 363      def visit_AugAssign(self, node):
 364  
 365          self.visit(node.target)
 366          self.write(' ' + BINOP_SYMBOLS[type(node.op)] + '= ')
 367          self.visit(node.value)
 368  
 369      def visit_ImportFrom(self, node):
 370  
 371          imports = []
 372          for alias in node.names:
 373              name = alias.name
 374              if alias.asname:
 375                  name += ' as ' + alias.asname
 376              imports.append(name)
 377          self.write('from {0}{1} import {2}'.format('.' * node.level, node.module or '', ', '.join(imports)))
 378  
 379      def visit_Import(self, node):
 380          write_comma = CommaWriter(self.write)
 381          self.write('import ')
 382          for item in node.names:
 383              write_comma()
 384              self.visit(item)
 385  
 386      def visit_Expr(self, node):
 387          self.correct_line_number(node)
 388          if isinstance(node.value, ast.Str):
 389              self.docstring(node.value)
 390          else:
 391              self.generic_visit(node)
 392  
 393      def visit_keyword(self, node):
 394          if self._is_node_args_valid(node, 'arg'):
 395              self.write(node.arg + '=')
 396          else:
 397              self.write('**')
 398          self.visit(node.value)
 399  
 400      def visit_FunctionDef(self, node):
 401          self.function_definition(node)
 402  
 403      def function_definition(self, node, prefixes=()):
 404          self.decorators(node)
 405  
 406          self._prefixes(prefixes)
 407          self.write('def %s(' % node.name)
 408          self.signature(node.args)
 409          self.write('):')
 410          self.body(node.body)
 411  
 412      def _prefixes(self, prefixes):
 413          self.write(' '.join(prefixes))
 414          if prefixes:
 415              self.write(' ')
 416  
 417      def visit_ClassDef(self, node):
 418          have_args = []
 419  
 420          def paren_or_comma():
 421              if have_args:
 422                  self.write(', ')
 423              else:
 424                  have_args.append(True)
 425                  self.write('(')
 426  
 427          self.decorators(node)
 428  
 429          self.write('class %s' % node.name)
 430          for base in node.bases:
 431              paren_or_comma()
 432              self.visit(base)
 433          self.write(have_args and '):' or ':')
 434          self.body(node.body)
 435  
 436      def visit_If(self, node):
 437          self.if_elif(node)
 438  
 439      def if_elif(self, node, use_elif=False):
 440          self.correct_line_number(node, within_statement=False)
 441          if use_elif:
 442              self.write('elif ')
 443          else:
 444              self.write('if ')
 445          self.visit(node.test)
 446          self.write(':')
 447          self.body(node.body)
 448          if node.orelse:
 449              self.if_or_else(node)
 450  
 451      def if_or_else(self, node):
 452          if len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If):
 453              self.if_elif(node.orelse[0], use_elif=True)
 454          else:
 455              self.or_else(node)
 456  
 457      def visit_For(self, node):
 458          self.for_loop(node)
 459  
 460      def for_loop(self, node, prefixes=()):
 461          self._prefixes(prefixes)
 462          self.write('for ')
 463          self.visit(node.target)
 464          self.write(' in ')
 465          self.visit(node.iter)
 466          self.write(':')
 467          self.body_or_else(node)
 468  
 469      def visit_While(self, node):
 470          self.write('while ')
 471          self.visit(node.test)
 472          self.write(':')
 473          self.body_or_else(node)
 474  
 475      def visit_Pass(self, node):
 476          self.write('pass')
 477  
 478      def visit_Print(self, node):
 479          self.correct_line_number(node)
 480          self.write('print ')
 481          want_comma = False
 482          if node.dest is not None:
 483              self.write('>> ')
 484              self.visit(node.dest)
 485              want_comma = True
 486          for value in node.values:
 487              if want_comma:
 488                  self.write(', ')
 489              self.visit(value)
 490              want_comma = True
 491          if not node.nl:
 492              self.write(',')
 493  
 494      def visit_Delete(self, node):
 495          self.write('del ')
 496          for target in node.targets:
 497              self.visit(target)
 498              if target is not node.targets[-1]:
 499                  self.write(', ')
 500  
 501      def visit_Global(self, node):
 502          self.write('global ' + ', '.join(node.names))
 503  
 504      def visit_Nonlocal(self, node):
 505          self.write('nonlocal ' + ', '.join(node.names))
 506  
 507      def visit_Return(self, node):
 508  
 509          self.write('return')
 510          if node.value:
 511              self.write(' ')
 512              self.visit(node.value)
 513  
 514      def visit_Break(self, node):
 515          self.write('break')
 516  
 517      def visit_Continue(self, node):
 518          self.write('continue')
 519  
 520      def visit_Raise(self, node):
 521  
 522          self.write('raise')
 523          if self._is_node_args_valid(node, 'exc'):
 524              self.raise_exc(node)
 525          elif self._is_node_args_valid(node, 'type'):
 526              self.raise_type(node)
 527  
 528      def raise_type(self, node):
 529          self.write(' ')
 530          self.visit(node.type)
 531          if node.inst is not None:
 532              self.write(', ')
 533              self.visit(node.inst)
 534          if node.tback is not None:
 535              self.write(', ')
 536              self.visit(node.tback)
 537  
 538      def raise_exc(self, node):
 539          self.write(' ')
 540          self.visit(node.exc)
 541          if node.cause is not None:
 542              self.write(' from ')
 543              self.visit(node.cause)
 544  
 545      # Expressions
 546  
 547      def visit_Attribute(self, node):
 548          self.visit(node.value)
 549          self.write('.' + node.attr)
 550  
 551      def visit_Call(self, node):
 552          self.visit(node.func)
 553          with self.inside('(', ')'):
 554              starargs = getattr(node, 'starargs', None)
 555              kwargs = getattr(node, 'kwargs', None)
 556              if starargs:
 557                  starargs = [starargs]
 558              else:
 559                  starargs = []
 560              if kwargs:
 561                  kwargs = [kwargs]
 562              else:
 563                  kwargs = []
 564              self.call_signature(node.args, node.keywords, starargs, kwargs)
 565  
 566      def call_signature(self, args, keywords, starargs, kwargs):
 567          write_comma = CommaWriter(self.write)
 568          self.call_signature_part(args, self.call_arg, write_comma)
 569          self.call_signature_part(keywords, self.call_keyword, write_comma)
 570          self.call_signature_part(starargs, self.call_starargs, write_comma)
 571          self.call_signature_part(kwargs, self.call_kwarg, write_comma)
 572  
 573      def call_signature_part(self, args, arg_processor, write_comma):
 574          for arg in args:
 575              write_comma()
 576              self.correct_line_number(arg, use_line_continuation=False)
 577              arg_processor(arg)
 578  
 579      def call_kwarg(self, kwarg):
 580          self.write('**')
 581          self.visit(kwarg)
 582  
 583      def call_starargs(self, stararg):
 584          self.write('*')
 585          self.visit(stararg)
 586  
 587      def call_keyword(self, keyword):
 588          self.visit(keyword)
 589  
 590      def call_arg(self, arg):
 591          self.visit(arg)
 592  
 593      def visit_Name(self, node):
 594          self.write(node.id)
 595  
 596      def visit_str(self, node):
 597          self.write(node)
 598  
 599      def visit_Str(self, node):
 600          self.write(repr(node.s))
 601  
 602      def visit_Bytes(self, node):
 603          self.write(repr(node.s))
 604  
 605      def visit_Num(self, node):
 606          value = node.n.imag if isinstance(node.n, complex) else node.n
 607  
 608          with self.inside('(', ')', cond=(value < 0)):
 609              self.write(repr(node.n))
 610  
 611      def visit_Tuple(self, node):
 612          with self.inside('(', ')'):
 613              idx = -1
 614              for idx, item in enumerate(node.elts):
 615                  if idx:
 616                      self.write(', ')
 617                  self.visit(item)
 618              if not idx:
 619                  self.write(',')
 620  
 621      def sequence_visit(left, right):  # @NoSelf
 622          def visit(self, node):
 623              with self.inside(left, right):
 624                  for idx, item in enumerate(node.elts):
 625                      if idx:
 626                          self.write(', ')
 627                      self.visit(item)
 628  
 629          return visit
 630  
 631      visit_List = sequence_visit('[', ']')
 632      visit_Set = sequence_visit('{', '}')
 633      del sequence_visit
 634  
 635      def visit_Dict(self, node):
 636          with self.inside('{', '}'):
 637              for idx, (key, value) in enumerate(zip(node.keys, node.values)):
 638                  if idx:
 639                      self.write(', ')
 640                  if key:
 641                      self.visit(key)
 642                      self.write(': ')
 643                  else:
 644                      self.write('**')
 645                  self.visit(value)
 646  
 647      def visit_BinOp(self, node):
 648          with self.inside('(', ')', cond=isinstance(node.parent, (ast.BinOp, ast.Attribute))):
 649              self.visit(node.left)
 650              self.write(' %s ' % BINOP_SYMBOLS[type(node.op)])
 651              self.visit(node.right)
 652  
 653      def visit_BoolOp(self, node):
 654          with self.inside('(', ')'):
 655              for idx, value in enumerate(node.values):
 656                  if idx:
 657                      self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)])
 658                  self.visit(value)
 659  
 660      def visit_Compare(self, node):
 661          with self.inside('(', ')', cond=(isinstance(node.parent, ast.Compare))):
 662              self.visit(node.left)
 663              for op, right in zip(node.ops, node.comparators):
 664                  self.write(' %s ' % CMPOP_SYMBOLS[type(op)])
 665                  self.visit(right)
 666  
 667      def visit_UnaryOp(self, node):
 668          with self.inside('(', ')', cond=isinstance(node.parent, (ast.BinOp, ast.UnaryOp))):
 669              op = UNARYOP_SYMBOLS[type(node.op)]
 670              self.write(op)
 671              if op == 'not':
 672                  self.write(' ')
 673  
 674              with self.inside('(', ')', cond=(not isinstance(node.operand, (ast.Name, ast.Num))
 675                                               and not self._is_named_constant(node.operand))):
 676                  self.visit(node.operand)
 677  
 678      def visit_Subscript(self, node):
 679          self.visit(node.value)
 680          with self.inside('[', ']'):
 681              self.visit(node.slice)
 682  
 683      def visit_Slice(self, node):
 684          self.slice_lower(node)
 685          self.write(':')
 686          self.slice_upper(node)
 687          self.slice_step(node)
 688  
 689      def slice_step(self, node):
 690          if node.step is not None:
 691              self.write(':')
 692              if not (isinstance(node.step, ast.Name) and node.step.id == 'None'):
 693                  self.visit(node.step)
 694  
 695      def slice_upper(self, node):
 696          if node.upper is not None:
 697              self.visit(node.upper)
 698  
 699      def slice_lower(self, node):
 700          if node.lower is not None:
 701              self.visit(node.lower)
 702  
 703      def visit_ExtSlice(self, node):
 704          for idx, item in enumerate(node.dims):
 705              if idx:
 706                  self.write(',')
 707              self.visit(item)
 708  
 709      def visit_Yield(self, node):
 710          self.write('yield')
 711          if node.value:
 712              self.write(' ')
 713              self.visit(node.value)
 714  
 715      def visit_Lambda(self, node):
 716          with self.inside('(', ')', cond=isinstance(node.parent, ast.Call)):
 717              self.write('lambda')
 718              self.signature(node.args, add_space=True)
 719              self.write(': ')
 720              with self.inside('(', ')'):
 721                  self.visit(node.body)
 722  
 723      def visit_Ellipsis(self, node):
 724          self.write('...')
 725  
 726      def generator_visit(left, right):  # @NoSelf
 727          def visit(self, node):
 728              self.write(left)
 729              self.visit(node.elt)
 730              for comprehension in node.generators:
 731                  self.visit(comprehension)
 732              self.write(right)
 733  
 734          return visit
 735  
 736      visit_ListComp = generator_visit('[', ']')
 737      visit_GeneratorExp = generator_visit('(', ')')
 738      visit_SetComp = generator_visit('{', '}')
 739      del generator_visit
 740  
 741      def visit_DictComp(self, node):
 742          with self.inside('{', '}'):
 743              self.visit(node.key)
 744              self.write(': ')
 745              self.visit(node.value)
 746              for comprehension in node.generators:
 747                  self.visit(comprehension)
 748  
 749      def visit_IfExp(self, node):
 750          with self.inside('(', ')', cond=isinstance(node.parent, ast.BinOp)):
 751              self.visit(node.body)
 752              self.write(' if ')
 753              self.visit(node.test)
 754              self.keyword_and_body(' else ', [node.orelse])
 755  
 756      def visit_Starred(self, node):
 757          self.write('*')
 758          self.visit(node.value)
 759  
 760      def visit_Repr(self, node):
 761          with self.inside('`', '`'):
 762              self.visit(node.value)
 763  
 764      # Helper Nodes
 765      def visit_alias(self, node):
 766          self.write(node.name)
 767          if node.asname is not None:
 768              self.write(' as ' + node.asname)
 769  
 770      def visit_comprehension(self, node):
 771          self.write(' for ')
 772          self.visit(node.target)
 773          self.write(' in ')
 774          self.visit(node.iter)
 775          if node.ifs:
 776              for if_ in node.ifs:
 777                  self.write(' if ')
 778                  self.visit(if_)
 779  
 780      def visit_ExceptHandler(self, node):
 781          self.write('except')
 782          if node.type is not None:
 783              self.write(' ')
 784              self.visit(node.type)
 785              if node.name is not None:
 786                  self.write(' as ')
 787                  self.visit(node.name)
 788          self.write(':')
 789          self.body(node.body)
 790  
 791      def visit_arg(self, node):
 792          self.write(node.arg)
 793  
 794      def visit_Assert(self, node):
 795          self.write('assert ')
 796          self.visit(node.test)
 797          if node.msg:
 798              self.write(', ')
 799              self.visit(node.msg)
 800  
 801      def visit_TryExcept(self, node):
 802          self.write('try:')
 803          self.body(node.body)
 804          if node.handlers:
 805              self.try_handlers(node)
 806          if node.orelse:
 807              self.or_else(node)
 808  
 809      def try_handlers(self, node):
 810          for handler in node.handlers:
 811              self.correct_line_number(handler, within_statement=False)
 812              self.visit(handler)
 813  
 814      def visit_TryFinally(self, node):
 815          self.write('try:')
 816          self.body(node.body)
 817          self.final_body(node)
 818  
 819      def final_body(self, node):
 820          self.keyword_and_body('finally:', node.finalbody)
 821  
 822      def visit_With(self, node):
 823          self.with_body(node)
 824  
 825      def with_body(self, node, prefixes=[]):
 826          self._prefixes(prefixes)
 827          self.write('with ')
 828          self.visit(node.context_expr)
 829          if node.optional_vars is not None:
 830              self.write(' as ')
 831              self.visit(node.optional_vars)
 832          self.write(':')
 833          self.body(node.body)
 834  
 835      @staticmethod
 836      def _is_named_constant(node):
 837          return isinstance(node, ast.Expr) and hasattr(node, 'value') and isinstance(node.value, ast.Name)
 838  
 839  
 840  class SourceGeneratorNodeVisitorPython26(BaseSourceGeneratorNodeVisitor):
 841      __python_version__ = (2, 6)
 842  
 843  
 844  class SourceGeneratorNodeVisitorPython27(SourceGeneratorNodeVisitorPython26):
 845      __python_version__ = (2, 7)
 846  
 847  
 848  class SourceGeneratorNodeVisitorPython30(SourceGeneratorNodeVisitorPython27):
 849      __python_version__ = (3, 0)
 850  
 851      def visit_ClassDef(self, node):
 852          have_args = []
 853  
 854          def paren_or_comma():
 855              if have_args:
 856                  self.write(', ')
 857              else:
 858                  have_args.append(True)
 859                  self.write('(')
 860  
 861          self.decorators(node)
 862          self.correct_line_number(node)
 863          self.write('class %s' % node.name)
 864          for base in node.bases:
 865              paren_or_comma()
 866              self.visit(base)
 867          if self._is_node_args_valid(node, 'keywords'):
 868              for keyword in node.keywords:
 869                  paren_or_comma()
 870                  self.visit(keyword)
 871          self.write(have_args and '):' or ':')
 872          self.body(node.body)
 873  
 874      def visit_FunctionDef(self, node):
 875          self.decorators(node)
 876  
 877          self.write('def %s(' % node.name)
 878          self.signature(node.args)
 879          self.write(')')
 880          if self._is_node_args_valid(node, 'returns'):
 881              self.write(' -> ')
 882              self.visit(node.returns)
 883          self.write(':')
 884          self.body(node.body)
 885  
 886  
 887  class SourceGeneratorNodeVisitorPython31(SourceGeneratorNodeVisitorPython30):
 888      __python_version__ = (3, 1)
 889  
 890  
 891  class SourceGeneratorNodeVisitorPython32(SourceGeneratorNodeVisitorPython31):
 892      __python_version__ = (3, 2)
 893  
 894  
 895  class SourceGeneratorNodeVisitorPython33(SourceGeneratorNodeVisitorPython32):
 896      __python_version__ = (3, 3)
 897  
 898      def visit_Try(self, node):
 899          self.write('try:')
 900          self.body(node.body)
 901          if node.handlers:
 902              self.try_handlers(node)
 903          if node.finalbody:
 904              self.final_body(node)
 905          if node.orelse:
 906              self.or_else(node)
 907  
 908      def with_body(self, node, prefixes=[]):
 909          self._prefixes(prefixes)
 910          self.write('with ')
 911          for with_item in node.items:
 912              self.visit(with_item.context_expr)
 913              if with_item.optional_vars is not None:
 914                  self.write(' as ')
 915                  self.visit(with_item.optional_vars)
 916              if with_item != node.items[-1]:
 917                  self.write(', ')
 918          self.write(':')
 919          self.body(node.body)
 920  
 921      def visit_YieldFrom(self, node):
 922          self.write('yield from ')
 923          self.visit(node.value)
 924  
 925  
 926  class SourceGeneratorNodeVisitorPython34(SourceGeneratorNodeVisitorPython33):
 927      __python_version__ = (3, 4)
 928  
 929      def visit_NameConstant(self, node):
 930          self.write(str(node.value))
 931  
 932      def visit_Name(self, node):
 933          if isinstance(node.id, ast.arg):
 934              self.write(node.id.arg)
 935          else:
 936              self.write(node.id)
 937  
 938      @staticmethod
 939      def _is_named_constant(node):
 940          return isinstance(node, ast.NameConstant)
 941  
 942  
 943  class SourceGeneratorNodeVisitorPython35(SourceGeneratorNodeVisitorPython34):
 944      __python_version__ = (3, 5)
 945  
 946      def visit_AsyncFunctionDef(self, node):
 947          self.function_definition(node, prefixes=['async'])
 948  
 949      def visit_AsyncFor(self, node):
 950          self.for_loop(node, prefixes=['async'])
 951  
 952      def visit_AsyncWith(self, node):
 953          self.with_body(node, prefixes=['async'])
 954  
 955      def visit_Await(self, node):
 956          self.write('await ')
 957          if self._is_node_args_valid(node, 'value'):
 958              self.visit(node.value)
 959  
 960      def visit_Call(self, node):
 961          self.visit(node.func)
 962          with self.inside('(', ')'):
 963              args, starargs = self._separate_args_and_starargs(node)
 964              keywords, kwargs = self._separate_keywords_and_kwargs(node)
 965              self.call_signature(args, keywords, starargs, kwargs)
 966  
 967      @staticmethod
 968      def _separate_keywords_and_kwargs(node):
 969          keywords = []
 970          kwargs = []
 971          for keyword in node.keywords:
 972              if keyword.arg:
 973                  keywords.append(keyword)
 974              else:
 975                  kwargs.append(keyword)
 976          return keywords, kwargs
 977  
 978      @staticmethod
 979      def _separate_args_and_starargs(node):
 980          args = []
 981          starargs = []
 982          for arg in node.args:
 983              if isinstance(arg, ast.Starred):
 984                  starargs.append(arg)
 985              else:
 986                  args.append(arg)
 987          return args, starargs
 988  
 989      def call_starargs(self, stararg):
 990          self.visit(stararg)
 991  
 992      def call_kwarg(self, kwarg):
 993          self.visit(kwarg)
 994  
 995  
 996  class SourceGeneratorNodeVisitorPython36(SourceGeneratorNodeVisitorPython35):
 997      __python_version__ = (3, 6)
 998  
 999      def visit_JoinedStr(self, node):
1000          if self._is_node_args_valid(node, 'values'):
1001              with self.inside('f\'', '\''):
1002                  for item in node.values:
1003                      if isinstance(item, ast.Str):
1004                          self.write(item.s.lstrip('\'').rstrip('\'').replace("'", "\\'"))
1005                      else:
1006                          self.visit(item)
1007  
1008      def visit_FormattedValue(self, node):
1009          if self._is_node_args_valid(node, 'value'):
1010              with self.inside('{', '}'):
1011                  self.visit(node.value)
1012                  if node.conversion != -1:
1013                      self.write('!%c' % (node.conversion,))
1014  
1015  
1016  class SourceGeneratorNodeVisitorPython38(SourceGeneratorNodeVisitorPython36):
1017      __python_version__ = (3, 8)
1018  
1019      def visit_Constant(self, node):
1020          if type(node.value) == str:
1021              self.write(repr(node.s))
1022          elif node.value == Ellipsis:
1023              self.write('...')
1024          else:
1025              self.write(str(node.value))
1026  
1027      def visit_NamedExpr(self, node):
1028          self.visit(node.target)
1029          self.write(' := ')
1030          self.visit(node.value)
1031  
1032      def signature(self, node, add_space=False):
1033          write_comma = CommaWriter(self.write, add_space_at_beginning=add_space)
1034  
1035  
1036          defaults = list(node.defaults)
1037  
1038          if node.posonlyargs:
1039              padding = [None] * (len(node.posonlyargs) - len(node.defaults))
1040              for arg, default in zip(node.posonlyargs, padding + defaults[:len(node.posonlyargs)]):
1041                  self.signature_arg(arg, default, write_comma)
1042              self.write(', /')
1043              defaults = defaults[len(node.posonlyargs):]
1044  
1045          padding = [None] * (len(node.args) - len(node.defaults))
1046          for arg, default in zip(node.args, padding + defaults):
1047              self.signature_arg(arg, default, write_comma)
1048  
1049          self.signature_spec_arg(node, 'vararg', write_comma, prefix='*')
1050          self.signature_kwonlyargs(node, write_comma)
1051          self.signature_spec_arg(node, 'kwarg', write_comma, prefix='**')
1052  
1053      @classmethod
1054      def _get_actual_lineno(cls, node):
1055          if isinstance(node, ast.FunctionDef) and node.decorator_list:
1056              return node.decorator_list[0].lineno
1057          else:
1058              return SourceGeneratorNodeVisitorPython36._get_actual_lineno(node)
1059  
1060  
1061  SourceGeneratorNodeVisitor = utils.get_by_python_version([
1062      SourceGeneratorNodeVisitorPython26,
1063      SourceGeneratorNodeVisitorPython27,
1064      SourceGeneratorNodeVisitorPython30,
1065      SourceGeneratorNodeVisitorPython31,
1066      SourceGeneratorNodeVisitorPython32,
1067      SourceGeneratorNodeVisitorPython33,
1068      SourceGeneratorNodeVisitorPython34,
1069      SourceGeneratorNodeVisitorPython35,
1070      SourceGeneratorNodeVisitorPython36,
1071      SourceGeneratorNodeVisitorPython38
1072  ])