Home | History | Annotate | Download | only in fixes 
      1 """Fixer for function definitions with tuple parameters.
      2 
      3 def func(((a, b), c), d):
      4     ...
      5 
      6     ->
      7 
      8 def func(x, d):
      9     ((a, b), c) = x
     10     ...
     11 
     12 It will also support lambdas:
     13 
     14     lambda (x, y): x + y -> lambda t: t[0] + t[1]
     15 
     16     # The parens are a syntax error in Python 3
     17     lambda (x): x + y -> lambda x: x + y
     18 """
     19 # Author: Collin Winter
     20 
     21 # Local imports
     22 from .. import pytree
     23 from ..pgen2 import token
     24 from .. import fixer_base
     25 from ..fixer_util import Assign, Name, Newline, Number, Subscript, syms
     26 
     27 def is_docstring(stmt):
     28     return isinstance(stmt, pytree.Node) and \
     29            stmt.children[0].type == token.STRING
     30 
     31 class FixTupleParams(fixer_base.BaseFix):
     32     run_order = 4 #use a lower order since lambda is part of other
     33                   #patterns
     34     BM_compatible = True
     35 
     36     PATTERN = """
     37               funcdef< 'def' any parameters< '(' args=any ')' >
     38                        ['->' any] ':' suite=any+ >
     39               |
     40               lambda=
     41               lambdef< 'lambda' args=vfpdef< '(' inner=any ')' >
     42                        ':' body=any
     43               >
     44               """
     45 
     46     def transform(self, node, results):
     47         if "lambda" in results:
     48             return self.transform_lambda(node, results)
     49 
     50         new_lines = []
     51         suite = results["suite"]
     52         args = results["args"]
     53         # This crap is so "def foo(...): x = 5; y = 7" is handled correctly.
     54         # TODO(cwinter): suite-cleanup
     55         if suite[0].children[1].type == token.INDENT:
     56             start = 2
     57             indent = suite[0].children[1].value
     58             end = Newline()
     59         else:
     60             start = 0
     61             indent = u"; "
     62             end = pytree.Leaf(token.INDENT, u"")
     63 
     64         # We need access to self for new_name(), and making this a method
     65         #  doesn't feel right. Closing over self and new_lines makes the
     66         #  code below cleaner.
     67         def handle_tuple(tuple_arg, add_prefix=False):
     68             n = Name(self.new_name())
     69             arg = tuple_arg.clone()
     70             arg.prefix = u""
     71             stmt = Assign(arg, n.clone())
     72             if add_prefix:
     73                 n.prefix = u" "
     74             tuple_arg.replace(n)
     75             new_lines.append(pytree.Node(syms.simple_stmt,
     76                                          [stmt, end.clone()]))
     77 
     78         if args.type == syms.tfpdef:
     79             handle_tuple(args)
     80         elif args.type == syms.typedargslist:
     81             for i, arg in enumerate(args.children):
     82                 if arg.type == syms.tfpdef:
     83                     # Without add_prefix, the emitted code is correct,
     84                     #  just ugly.
     85                     handle_tuple(arg, add_prefix=(i > 0))
     86 
     87         if not new_lines:
     88             return
     89 
     90         # This isn't strictly necessary, but it plays nicely with other fixers.
     91         # TODO(cwinter) get rid of this when children becomes a smart list
     92         for line in new_lines:
     93             line.parent = suite[0]
     94 
     95         # TODO(cwinter) suite-cleanup
     96         after = start
     97         if start == 0:
     98             new_lines[0].prefix = u" "
     99         elif is_docstring(suite[0].children[start]):
    100             new_lines[0].prefix = indent
    101             after = start + 1
    102 
    103         for line in new_lines:
    104             line.parent = suite[0]
    105         suite[0].children[after:after] = new_lines
    106         for i in range(after+1, after+len(new_lines)+1):
    107             suite[0].children[i].prefix = indent
    108         suite[0].changed()
    109 
    110     def transform_lambda(self, node, results):
    111         args = results["args"]
    112         body = results["body"]
    113         inner = simplify_args(results["inner"])
    114 
    115         # Replace lambda ((((x)))): x  with lambda x: x
    116         if inner.type == token.NAME:
    117             inner = inner.clone()
    118             inner.prefix = u" "
    119             args.replace(inner)
    120             return
    121 
    122         params = find_params(args)
    123         to_index = map_to_index(params)
    124         tup_name = self.new_name(tuple_name(params))
    125 
    126         new_param = Name(tup_name, prefix=u" ")
    127         args.replace(new_param.clone())
    128         for n in body.post_order():
    129             if n.type == token.NAME and n.value in to_index:
    130                 subscripts = [c.clone() for c in to_index[n.value]]
    131                 new = pytree.Node(syms.power,
    132                                   [new_param.clone()] + subscripts)
    133                 new.prefix = n.prefix
    134                 n.replace(new)
    135 
    136 
    137 ### Helper functions for transform_lambda()
    138 
    139 def simplify_args(node):
    140     if node.type in (syms.vfplist, token.NAME):
    141         return node
    142     elif node.type == syms.vfpdef:
    143         # These look like vfpdef< '(' x ')' > where x is NAME
    144         # or another vfpdef instance (leading to recursion).
    145         while node.type == syms.vfpdef:
    146             node = node.children[1]
    147         return node
    148     raise RuntimeError("Received unexpected node %s" % node)
    149 
    150 def find_params(node):
    151     if node.type == syms.vfpdef:
    152         return find_params(node.children[1])
    153     elif node.type == token.NAME:
    154         return node.value
    155     return [find_params(c) for c in node.children if c.type != token.COMMA]
    156 
    157 def map_to_index(param_list, prefix=[], d=None):
    158     if d is None:
    159         d = {}
    160     for i, obj in enumerate(param_list):
    161         trailer = [Subscript(Number(unicode(i)))]
    162         if isinstance(obj, list):
    163             map_to_index(obj, trailer, d=d)
    164         else:
    165             d[obj] = prefix + trailer
    166     return d
    167 
    168 def tuple_name(param_list):
    169     l = []
    170     for obj in param_list:
    171         if isinstance(obj, list):
    172             l.append(tuple_name(obj))
    173         else:
    174             l.append(obj)
    175     return u"_".join(l)
    176