fix_idioms.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. """Adjust some old Python 2 idioms to their modern counterparts.
  2. * Change some type comparisons to isinstance() calls:
  3. type(x) == T -> isinstance(x, T)
  4. type(x) is T -> isinstance(x, T)
  5. type(x) != T -> not isinstance(x, T)
  6. type(x) is not T -> not isinstance(x, T)
  7. * Change "while 1:" into "while True:".
  8. * Change both
  9. v = list(EXPR)
  10. v.sort()
  11. foo(v)
  12. and the more general
  13. v = EXPR
  14. v.sort()
  15. foo(v)
  16. into
  17. v = sorted(EXPR)
  18. foo(v)
  19. """
  20. # Author: Jacques Frechet, Collin Winter
  21. # Local imports
  22. from .. import fixer_base
  23. from ..fixer_util import Call, Comma, Name, Node, BlankLine, syms
  24. CMP = "(n='!=' | '==' | 'is' | n=comp_op< 'is' 'not' >)"
  25. TYPE = "power< 'type' trailer< '(' x=any ')' > >"
  26. class FixIdioms(fixer_base.BaseFix):
  27. explicit = True # The user must ask for this fixer
  28. PATTERN = r"""
  29. isinstance=comparison< %s %s T=any >
  30. |
  31. isinstance=comparison< T=any %s %s >
  32. |
  33. while_stmt< 'while' while='1' ':' any+ >
  34. |
  35. sorted=any<
  36. any*
  37. simple_stmt<
  38. expr_stmt< id1=any '='
  39. power< list='list' trailer< '(' (not arglist<any+>) any ')' > >
  40. >
  41. '\n'
  42. >
  43. sort=
  44. simple_stmt<
  45. power< id2=any
  46. trailer< '.' 'sort' > trailer< '(' ')' >
  47. >
  48. '\n'
  49. >
  50. next=any*
  51. >
  52. |
  53. sorted=any<
  54. any*
  55. simple_stmt< expr_stmt< id1=any '=' expr=any > '\n' >
  56. sort=
  57. simple_stmt<
  58. power< id2=any
  59. trailer< '.' 'sort' > trailer< '(' ')' >
  60. >
  61. '\n'
  62. >
  63. next=any*
  64. >
  65. """ % (TYPE, CMP, CMP, TYPE)
  66. def match(self, node):
  67. r = super(FixIdioms, self).match(node)
  68. # If we've matched one of the sort/sorted subpatterns above, we
  69. # want to reject matches where the initial assignment and the
  70. # subsequent .sort() call involve different identifiers.
  71. if r and "sorted" in r:
  72. if r["id1"] == r["id2"]:
  73. return r
  74. return None
  75. return r
  76. def transform(self, node, results):
  77. if "isinstance" in results:
  78. return self.transform_isinstance(node, results)
  79. elif "while" in results:
  80. return self.transform_while(node, results)
  81. elif "sorted" in results:
  82. return self.transform_sort(node, results)
  83. else:
  84. raise RuntimeError("Invalid match")
  85. def transform_isinstance(self, node, results):
  86. x = results["x"].clone() # The thing inside of type()
  87. T = results["T"].clone() # The type being compared against
  88. x.prefix = ""
  89. T.prefix = " "
  90. test = Call(Name("isinstance"), [x, Comma(), T])
  91. if "n" in results:
  92. test.prefix = " "
  93. test = Node(syms.not_test, [Name("not"), test])
  94. test.prefix = node.prefix
  95. return test
  96. def transform_while(self, node, results):
  97. one = results["while"]
  98. one.replace(Name("True", prefix=one.prefix))
  99. def transform_sort(self, node, results):
  100. sort_stmt = results["sort"]
  101. next_stmt = results["next"]
  102. list_call = results.get("list")
  103. simple_expr = results.get("expr")
  104. if list_call:
  105. list_call.replace(Name("sorted", prefix=list_call.prefix))
  106. elif simple_expr:
  107. new = simple_expr.clone()
  108. new.prefix = ""
  109. simple_expr.replace(Call(Name("sorted"), [new],
  110. prefix=simple_expr.prefix))
  111. else:
  112. raise RuntimeError("should not have reached here")
  113. sort_stmt.remove()
  114. btwn = sort_stmt.prefix
  115. # Keep any prefix lines between the sort_stmt and the list_call and
  116. # shove them right after the sorted() call.
  117. if "\n" in btwn:
  118. if next_stmt:
  119. # The new prefix should be everything from the sort_stmt's
  120. # prefix up to the last newline, then the old prefix after a new
  121. # line.
  122. prefix_lines = (btwn.rpartition("\n")[0], next_stmt[0].prefix)
  123. next_stmt[0].prefix = "\n".join(prefix_lines)
  124. else:
  125. assert list_call.parent
  126. assert list_call.next_sibling is None
  127. # Put a blank line after list_call and set its prefix.
  128. end_line = BlankLine()
  129. list_call.parent.append_child(end_line)
  130. assert list_call.next_sibling is end_line
  131. # The new prefix should be everything up to the first new line
  132. # of sort_stmt's prefix.
  133. end_line.prefix = btwn.rpartition("\n")[0]