Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame^] | 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 basefix |
| 25 | from .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(basefix.BaseFix): |
| 32 | PATTERN = """ |
| 33 | funcdef< 'def' any parameters< '(' args=any ')' > |
| 34 | ['->' any] ':' suite=any+ > |
| 35 | | |
| 36 | lambda= |
| 37 | lambdef< 'lambda' args=vfpdef< '(' inner=any ')' > |
| 38 | ':' body=any |
| 39 | > |
| 40 | """ |
| 41 | |
| 42 | def transform(self, node, results): |
| 43 | if "lambda" in results: |
| 44 | return self.transform_lambda(node, results) |
| 45 | |
| 46 | new_lines = [] |
| 47 | suite = results["suite"] |
| 48 | args = results["args"] |
| 49 | # This crap is so "def foo(...): x = 5; y = 7" is handled correctly. |
| 50 | # TODO(cwinter): suite-cleanup |
| 51 | if suite[0].children[1].type == token.INDENT: |
| 52 | start = 2 |
| 53 | indent = suite[0].children[1].value |
| 54 | end = Newline() |
| 55 | else: |
| 56 | start = 0 |
| 57 | indent = "; " |
| 58 | end = pytree.Leaf(token.INDENT, "") |
| 59 | |
| 60 | # We need access to self for new_name(), and making this a method |
| 61 | # doesn't feel right. Closing over self and new_lines makes the |
| 62 | # code below cleaner. |
| 63 | def handle_tuple(tuple_arg, add_prefix=False): |
| 64 | n = Name(self.new_name()) |
| 65 | arg = tuple_arg.clone() |
| 66 | arg.set_prefix("") |
| 67 | stmt = Assign(arg, n.clone()) |
| 68 | if add_prefix: |
| 69 | n.set_prefix(" ") |
| 70 | tuple_arg.replace(n) |
| 71 | new_lines.append(pytree.Node(syms.simple_stmt, |
| 72 | [stmt, end.clone()])) |
| 73 | |
| 74 | if args.type == syms.tfpdef: |
| 75 | handle_tuple(args) |
| 76 | elif args.type == syms.typedargslist: |
| 77 | for i, arg in enumerate(args.children): |
| 78 | if arg.type == syms.tfpdef: |
| 79 | # Without add_prefix, the emitted code is correct, |
| 80 | # just ugly. |
| 81 | handle_tuple(arg, add_prefix=(i > 0)) |
| 82 | |
| 83 | if not new_lines: |
| 84 | return node |
| 85 | |
| 86 | # This isn't strictly necessary, but it plays nicely with other fixers. |
| 87 | # TODO(cwinter) get rid of this when children becomes a smart list |
| 88 | for line in new_lines: |
| 89 | line.parent = suite[0] |
| 90 | |
| 91 | # TODO(cwinter) suite-cleanup |
| 92 | after = start |
| 93 | if start == 0: |
| 94 | new_lines[0].set_prefix(" ") |
| 95 | elif is_docstring(suite[0].children[start]): |
| 96 | new_lines[0].set_prefix(indent) |
| 97 | after = start + 1 |
| 98 | |
| 99 | suite[0].children[after:after] = new_lines |
| 100 | for i in range(after+1, after+len(new_lines)+1): |
| 101 | suite[0].children[i].set_prefix(indent) |
| 102 | suite[0].changed() |
| 103 | |
| 104 | def transform_lambda(self, node, results): |
| 105 | args = results["args"] |
| 106 | body = results["body"] |
| 107 | inner = simplify_args(results["inner"]) |
| 108 | |
| 109 | # Replace lambda ((((x)))): x with lambda x: x |
| 110 | if inner.type == token.NAME: |
| 111 | inner = inner.clone() |
| 112 | inner.set_prefix(" ") |
| 113 | args.replace(inner) |
| 114 | return |
| 115 | |
| 116 | params = find_params(args) |
| 117 | to_index = map_to_index(params) |
| 118 | tup_name = self.new_name(tuple_name(params)) |
| 119 | |
| 120 | new_param = Name(tup_name, prefix=" ") |
| 121 | args.replace(new_param.clone()) |
| 122 | for n in body.post_order(): |
| 123 | if n.type == token.NAME and n.value in to_index: |
| 124 | subscripts = [c.clone() for c in to_index[n.value]] |
| 125 | new = pytree.Node(syms.power, |
| 126 | [new_param.clone()] + subscripts) |
| 127 | new.set_prefix(n.get_prefix()) |
| 128 | n.replace(new) |
| 129 | |
| 130 | |
| 131 | ### Helper functions for transform_lambda() |
| 132 | |
| 133 | def simplify_args(node): |
| 134 | if node.type in (syms.vfplist, token.NAME): |
| 135 | return node |
| 136 | elif node.type == syms.vfpdef: |
| 137 | # These look like vfpdef< '(' x ')' > where x is NAME |
| 138 | # or another vfpdef instance (leading to recursion). |
| 139 | while node.type == syms.vfpdef: |
| 140 | node = node.children[1] |
| 141 | return node |
| 142 | raise RuntimeError("Received unexpected node %s" % node) |
| 143 | |
| 144 | def find_params(node): |
| 145 | if node.type == syms.vfpdef: |
| 146 | return find_params(node.children[1]) |
| 147 | elif node.type == token.NAME: |
| 148 | return node.value |
| 149 | return [find_params(c) for c in node.children if c.type != token.COMMA] |
| 150 | |
| 151 | def map_to_index(param_list, prefix=[], d=None): |
| 152 | if d is None: |
| 153 | d = {} |
| 154 | for i, obj in enumerate(param_list): |
| 155 | trailer = [Subscript(Number(i))] |
| 156 | if isinstance(obj, list): |
| 157 | map_to_index(obj, trailer, d=d) |
| 158 | else: |
| 159 | d[obj] = prefix + trailer |
| 160 | return d |
| 161 | |
| 162 | def tuple_name(param_list): |
| 163 | l = [] |
| 164 | for obj in param_list: |
| 165 | if isinstance(obj, list): |
| 166 | l.append(tuple_name(obj)) |
| 167 | else: |
| 168 | l.append(obj) |
| 169 | return "_".join(l) |