fixer_util.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520
  1. """
  2. Utility functions from 2to3, 3to2 and python-modernize (and some home-grown
  3. ones).
  4. Licences:
  5. 2to3: PSF License v2
  6. 3to2: Apache Software License (from 3to2/setup.py)
  7. python-modernize licence: BSD (from python-modernize/LICENSE)
  8. """
  9. from lib2to3.fixer_util import (FromImport, Newline, is_import,
  10. find_root, does_tree_import, Comma)
  11. from lib2to3.pytree import Leaf, Node
  12. from lib2to3.pygram import python_symbols as syms, python_grammar
  13. from lib2to3.pygram import token
  14. from lib2to3.fixer_util import (Node, Call, Name, syms, Comma, Number)
  15. import re
  16. def canonical_fix_name(fix, avail_fixes):
  17. """
  18. Examples:
  19. >>> canonical_fix_name('fix_wrap_text_literals')
  20. 'libfuturize.fixes.fix_wrap_text_literals'
  21. >>> canonical_fix_name('wrap_text_literals')
  22. 'libfuturize.fixes.fix_wrap_text_literals'
  23. >>> canonical_fix_name('wrap_te')
  24. ValueError("unknown fixer name")
  25. >>> canonical_fix_name('wrap')
  26. ValueError("ambiguous fixer name")
  27. """
  28. if ".fix_" in fix:
  29. return fix
  30. else:
  31. if fix.startswith('fix_'):
  32. fix = fix[4:]
  33. # Infer the full module name for the fixer.
  34. # First ensure that no names clash (e.g.
  35. # lib2to3.fixes.fix_blah and libfuturize.fixes.fix_blah):
  36. found = [f for f in avail_fixes
  37. if f.endswith('fix_{0}'.format(fix))]
  38. if len(found) > 1:
  39. raise ValueError("Ambiguous fixer name. Choose a fully qualified "
  40. "module name instead from these:\n" +
  41. "\n".join(" " + myf for myf in found))
  42. elif len(found) == 0:
  43. raise ValueError("Unknown fixer. Use --list-fixes or -l for a list.")
  44. return found[0]
  45. ## These functions are from 3to2 by Joe Amenta:
  46. def Star(prefix=None):
  47. return Leaf(token.STAR, u'*', prefix=prefix)
  48. def DoubleStar(prefix=None):
  49. return Leaf(token.DOUBLESTAR, u'**', prefix=prefix)
  50. def Minus(prefix=None):
  51. return Leaf(token.MINUS, u'-', prefix=prefix)
  52. def commatize(leafs):
  53. """
  54. Accepts/turns: (Name, Name, ..., Name, Name)
  55. Returns/into: (Name, Comma, Name, Comma, ..., Name, Comma, Name)
  56. """
  57. new_leafs = []
  58. for leaf in leafs:
  59. new_leafs.append(leaf)
  60. new_leafs.append(Comma())
  61. del new_leafs[-1]
  62. return new_leafs
  63. def indentation(node):
  64. """
  65. Returns the indentation for this node
  66. Iff a node is in a suite, then it has indentation.
  67. """
  68. while node.parent is not None and node.parent.type != syms.suite:
  69. node = node.parent
  70. if node.parent is None:
  71. return u""
  72. # The first three children of a suite are NEWLINE, INDENT, (some other node)
  73. # INDENT.value contains the indentation for this suite
  74. # anything after (some other node) has the indentation as its prefix.
  75. if node.type == token.INDENT:
  76. return node.value
  77. elif node.prev_sibling is not None and node.prev_sibling.type == token.INDENT:
  78. return node.prev_sibling.value
  79. elif node.prev_sibling is None:
  80. return u""
  81. else:
  82. return node.prefix
  83. def indentation_step(node):
  84. """
  85. Dirty little trick to get the difference between each indentation level
  86. Implemented by finding the shortest indentation string
  87. (technically, the "least" of all of the indentation strings, but
  88. tabs and spaces mixed won't get this far, so those are synonymous.)
  89. """
  90. r = find_root(node)
  91. # Collect all indentations into one set.
  92. all_indents = set(i.value for i in r.pre_order() if i.type == token.INDENT)
  93. if not all_indents:
  94. # nothing is indented anywhere, so we get to pick what we want
  95. return u" " # four spaces is a popular convention
  96. else:
  97. return min(all_indents)
  98. def suitify(parent):
  99. """
  100. Turn the stuff after the first colon in parent's children
  101. into a suite, if it wasn't already
  102. """
  103. for node in parent.children:
  104. if node.type == syms.suite:
  105. # already in the prefered format, do nothing
  106. return
  107. # One-liners have no suite node, we have to fake one up
  108. for i, node in enumerate(parent.children):
  109. if node.type == token.COLON:
  110. break
  111. else:
  112. raise ValueError(u"No class suite and no ':'!")
  113. # Move everything into a suite node
  114. suite = Node(syms.suite, [Newline(), Leaf(token.INDENT, indentation(node) + indentation_step(node))])
  115. one_node = parent.children[i+1]
  116. one_node.remove()
  117. one_node.prefix = u''
  118. suite.append_child(one_node)
  119. parent.append_child(suite)
  120. def NameImport(package, as_name=None, prefix=None):
  121. """
  122. Accepts a package (Name node), name to import it as (string), and
  123. optional prefix and returns a node:
  124. import <package> [as <as_name>]
  125. """
  126. if prefix is None:
  127. prefix = u""
  128. children = [Name(u"import", prefix=prefix), package]
  129. if as_name is not None:
  130. children.extend([Name(u"as", prefix=u" "),
  131. Name(as_name, prefix=u" ")])
  132. return Node(syms.import_name, children)
  133. _compound_stmts = (syms.if_stmt, syms.while_stmt, syms.for_stmt, syms.try_stmt, syms.with_stmt)
  134. _import_stmts = (syms.import_name, syms.import_from)
  135. def import_binding_scope(node):
  136. """
  137. Generator yields all nodes for which a node (an import_stmt) has scope
  138. The purpose of this is for a call to _find() on each of them
  139. """
  140. # import_name / import_from are small_stmts
  141. assert node.type in _import_stmts
  142. test = node.next_sibling
  143. # A small_stmt can only be followed by a SEMI or a NEWLINE.
  144. while test.type == token.SEMI:
  145. nxt = test.next_sibling
  146. # A SEMI can only be followed by a small_stmt or a NEWLINE
  147. if nxt.type == token.NEWLINE:
  148. break
  149. else:
  150. yield nxt
  151. # A small_stmt can only be followed by either a SEMI or a NEWLINE
  152. test = nxt.next_sibling
  153. # Covered all subsequent small_stmts after the import_stmt
  154. # Now to cover all subsequent stmts after the parent simple_stmt
  155. parent = node.parent
  156. assert parent.type == syms.simple_stmt
  157. test = parent.next_sibling
  158. while test is not None:
  159. # Yes, this will yield NEWLINE and DEDENT. Deal with it.
  160. yield test
  161. test = test.next_sibling
  162. context = parent.parent
  163. # Recursively yield nodes following imports inside of a if/while/for/try/with statement
  164. if context.type in _compound_stmts:
  165. # import is in a one-liner
  166. c = context
  167. while c.next_sibling is not None:
  168. yield c.next_sibling
  169. c = c.next_sibling
  170. context = context.parent
  171. # Can't chain one-liners on one line, so that takes care of that.
  172. p = context.parent
  173. if p is None:
  174. return
  175. # in a multi-line suite
  176. while p.type in _compound_stmts:
  177. if context.type == syms.suite:
  178. yield context
  179. context = context.next_sibling
  180. if context is None:
  181. context = p.parent
  182. p = context.parent
  183. if p is None:
  184. break
  185. def ImportAsName(name, as_name, prefix=None):
  186. new_name = Name(name)
  187. new_as = Name(u"as", prefix=u" ")
  188. new_as_name = Name(as_name, prefix=u" ")
  189. new_node = Node(syms.import_as_name, [new_name, new_as, new_as_name])
  190. if prefix is not None:
  191. new_node.prefix = prefix
  192. return new_node
  193. def is_docstring(node):
  194. """
  195. Returns True if the node appears to be a docstring
  196. """
  197. return (node.type == syms.simple_stmt and
  198. len(node.children) > 0 and node.children[0].type == token.STRING)
  199. def future_import(feature, node):
  200. """
  201. This seems to work
  202. """
  203. root = find_root(node)
  204. if does_tree_import(u"__future__", feature, node):
  205. return
  206. # Look for a shebang or encoding line
  207. shebang_encoding_idx = None
  208. for idx, node in enumerate(root.children):
  209. # Is it a shebang or encoding line?
  210. if is_shebang_comment(node) or is_encoding_comment(node):
  211. shebang_encoding_idx = idx
  212. if is_docstring(node):
  213. # skip over docstring
  214. continue
  215. names = check_future_import(node)
  216. if not names:
  217. # not a future statement; need to insert before this
  218. break
  219. if feature in names:
  220. # already imported
  221. return
  222. import_ = FromImport(u'__future__', [Leaf(token.NAME, feature, prefix=" ")])
  223. if shebang_encoding_idx == 0 and idx == 0:
  224. # If this __future__ import would go on the first line,
  225. # detach the shebang / encoding prefix from the current first line.
  226. # and attach it to our new __future__ import node.
  227. import_.prefix = root.children[0].prefix
  228. root.children[0].prefix = u''
  229. # End the __future__ import line with a newline and add a blank line
  230. # afterwards:
  231. children = [import_ , Newline()]
  232. root.insert_child(idx, Node(syms.simple_stmt, children))
  233. def future_import2(feature, node):
  234. """
  235. An alternative to future_import() which might not work ...
  236. """
  237. root = find_root(node)
  238. if does_tree_import(u"__future__", feature, node):
  239. return
  240. insert_pos = 0
  241. for idx, node in enumerate(root.children):
  242. if node.type == syms.simple_stmt and node.children and \
  243. node.children[0].type == token.STRING:
  244. insert_pos = idx + 1
  245. break
  246. for thing_after in root.children[insert_pos:]:
  247. if thing_after.type == token.NEWLINE:
  248. insert_pos += 1
  249. continue
  250. prefix = thing_after.prefix
  251. thing_after.prefix = u""
  252. break
  253. else:
  254. prefix = u""
  255. import_ = FromImport(u"__future__", [Leaf(token.NAME, feature, prefix=u" ")])
  256. children = [import_, Newline()]
  257. root.insert_child(insert_pos, Node(syms.simple_stmt, children, prefix=prefix))
  258. def parse_args(arglist, scheme):
  259. u"""
  260. Parse a list of arguments into a dict
  261. """
  262. arglist = [i for i in arglist if i.type != token.COMMA]
  263. ret_mapping = dict([(k, None) for k in scheme])
  264. for i, arg in enumerate(arglist):
  265. if arg.type == syms.argument and arg.children[1].type == token.EQUAL:
  266. # argument < NAME '=' any >
  267. slot = arg.children[0].value
  268. ret_mapping[slot] = arg.children[2]
  269. else:
  270. slot = scheme[i]
  271. ret_mapping[slot] = arg
  272. return ret_mapping
  273. # def is_import_from(node):
  274. # """Returns true if the node is a statement "from ... import ..."
  275. # """
  276. # return node.type == syms.import_from
  277. def is_import_stmt(node):
  278. return (node.type == syms.simple_stmt and node.children and
  279. is_import(node.children[0]))
  280. def touch_import_top(package, name_to_import, node):
  281. """Works like `does_tree_import` but adds an import statement at the
  282. top if it was not imported (but below any __future__ imports) and below any
  283. comments such as shebang lines).
  284. Based on lib2to3.fixer_util.touch_import()
  285. Calling this multiple times adds the imports in reverse order.
  286. Also adds "standard_library.install_aliases()" after "from future import
  287. standard_library". This should probably be factored into another function.
  288. """
  289. root = find_root(node)
  290. if does_tree_import(package, name_to_import, root):
  291. return
  292. # Ideally, we would look for whether futurize --all-imports has been run,
  293. # as indicated by the presence of ``from builtins import (ascii, ...,
  294. # zip)`` -- and, if it has, we wouldn't import the name again.
  295. # Look for __future__ imports and insert below them
  296. found = False
  297. for name in ['absolute_import', 'division', 'print_function',
  298. 'unicode_literals']:
  299. if does_tree_import('__future__', name, root):
  300. found = True
  301. break
  302. if found:
  303. # At least one __future__ import. We want to loop until we've seen them
  304. # all.
  305. start, end = None, None
  306. for idx, node in enumerate(root.children):
  307. if check_future_import(node):
  308. start = idx
  309. # Start looping
  310. idx2 = start
  311. while node:
  312. node = node.next_sibling
  313. idx2 += 1
  314. if not check_future_import(node):
  315. end = idx2
  316. break
  317. break
  318. assert start is not None
  319. assert end is not None
  320. insert_pos = end
  321. else:
  322. # No __future__ imports.
  323. # We look for a docstring and insert the new node below that. If no docstring
  324. # exists, just insert the node at the top.
  325. for idx, node in enumerate(root.children):
  326. if node.type != syms.simple_stmt:
  327. break
  328. if not is_docstring(node):
  329. # This is the usual case.
  330. break
  331. insert_pos = idx
  332. if package is None:
  333. import_ = Node(syms.import_name, [
  334. Leaf(token.NAME, u"import"),
  335. Leaf(token.NAME, name_to_import, prefix=u" ")
  336. ])
  337. else:
  338. import_ = FromImport(package, [Leaf(token.NAME, name_to_import, prefix=u" ")])
  339. if name_to_import == u'standard_library':
  340. # Add:
  341. # standard_library.install_aliases()
  342. # after:
  343. # from future import standard_library
  344. install_hooks = Node(syms.simple_stmt,
  345. [Node(syms.power,
  346. [Leaf(token.NAME, u'standard_library'),
  347. Node(syms.trailer, [Leaf(token.DOT, u'.'),
  348. Leaf(token.NAME, u'install_aliases')]),
  349. Node(syms.trailer, [Leaf(token.LPAR, u'('),
  350. Leaf(token.RPAR, u')')])
  351. ])
  352. ]
  353. )
  354. children_hooks = [install_hooks, Newline()]
  355. else:
  356. children_hooks = []
  357. # FromImport(package, [Leaf(token.NAME, name_to_import, prefix=u" ")])
  358. children_import = [import_, Newline()]
  359. old_prefix = root.children[insert_pos].prefix
  360. root.children[insert_pos].prefix = u''
  361. root.insert_child(insert_pos, Node(syms.simple_stmt, children_import, prefix=old_prefix))
  362. if len(children_hooks) > 0:
  363. root.insert_child(insert_pos + 1, Node(syms.simple_stmt, children_hooks))
  364. ## The following functions are from python-modernize by Armin Ronacher:
  365. # (a little edited).
  366. def check_future_import(node):
  367. """If this is a future import, return set of symbols that are imported,
  368. else return None."""
  369. # node should be the import statement here
  370. savenode = node
  371. if not (node.type == syms.simple_stmt and node.children):
  372. return set()
  373. node = node.children[0]
  374. # now node is the import_from node
  375. if not (node.type == syms.import_from and
  376. # node.type == token.NAME and # seems to break it
  377. hasattr(node.children[1], 'value') and
  378. node.children[1].value == u'__future__'):
  379. return set()
  380. if node.children[3].type == token.LPAR:
  381. node = node.children[4]
  382. else:
  383. node = node.children[3]
  384. # now node is the import_as_name[s]
  385. # print(python_grammar.number2symbol[node.type]) # breaks sometimes
  386. if node.type == syms.import_as_names:
  387. result = set()
  388. for n in node.children:
  389. if n.type == token.NAME:
  390. result.add(n.value)
  391. elif n.type == syms.import_as_name:
  392. n = n.children[0]
  393. assert n.type == token.NAME
  394. result.add(n.value)
  395. return result
  396. elif node.type == syms.import_as_name:
  397. node = node.children[0]
  398. assert node.type == token.NAME
  399. return set([node.value])
  400. elif node.type == token.NAME:
  401. return set([node.value])
  402. else:
  403. # TODO: handle brackets like this:
  404. # from __future__ import (absolute_import, division)
  405. assert False, "strange import: %s" % savenode
  406. SHEBANG_REGEX = r'^#!.*python'
  407. ENCODING_REGEX = r"^#.*coding[:=]\s*([-\w.]+)"
  408. def is_shebang_comment(node):
  409. """
  410. Comments are prefixes for Leaf nodes. Returns whether the given node has a
  411. prefix that looks like a shebang line or an encoding line:
  412. #!/usr/bin/env python
  413. #!/usr/bin/python3
  414. """
  415. return bool(re.match(SHEBANG_REGEX, node.prefix))
  416. def is_encoding_comment(node):
  417. """
  418. Comments are prefixes for Leaf nodes. Returns whether the given node has a
  419. prefix that looks like an encoding line:
  420. # coding: utf-8
  421. # encoding: utf-8
  422. # -*- coding: <encoding name> -*-
  423. # vim: set fileencoding=<encoding name> :
  424. """
  425. return bool(re.match(ENCODING_REGEX, node.prefix))
  426. def wrap_in_fn_call(fn_name, args, prefix=None):
  427. """
  428. Example:
  429. >>> wrap_in_fn_call("oldstr", (arg,))
  430. oldstr(arg)
  431. >>> wrap_in_fn_call("olddiv", (arg1, arg2))
  432. olddiv(arg1, arg2)
  433. >>> wrap_in_fn_call("olddiv", [arg1, comma, arg2, comma, arg3])
  434. olddiv(arg1, arg2, arg3)
  435. """
  436. assert len(args) > 0
  437. if len(args) == 2:
  438. expr1, expr2 = args
  439. newargs = [expr1, Comma(), expr2]
  440. else:
  441. newargs = args
  442. return Call(Name(fn_name), newargs, prefix=prefix)