fix_dict.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  1. # Copyright 2007 Google, Inc. All Rights Reserved.
  2. # Licensed to PSF under a Contributor Agreement.
  3. """Fixer for dict methods.
  4. d.keys() -> list(d.keys())
  5. d.items() -> list(d.items())
  6. d.values() -> list(d.values())
  7. d.iterkeys() -> iter(d.keys())
  8. d.iteritems() -> iter(d.items())
  9. d.itervalues() -> iter(d.values())
  10. d.viewkeys() -> d.keys()
  11. d.viewitems() -> d.items()
  12. d.viewvalues() -> d.values()
  13. Except in certain very specific contexts: the iter() can be dropped
  14. when the context is list(), sorted(), iter() or for...in; the list()
  15. can be dropped when the context is list() or sorted() (but not iter()
  16. or for...in!). Special contexts that apply to both: list(), sorted(), tuple()
  17. set(), any(), all(), sum().
  18. Note: iter(d.keys()) could be written as iter(d) but since the
  19. original d.iterkeys() was also redundant we don't fix this. And there
  20. are (rare) contexts where it makes a difference (e.g. when passing it
  21. as an argument to a function that introspects the argument).
  22. """
  23. # Local imports
  24. from .. import pytree
  25. from .. import patcomp
  26. from .. import fixer_base
  27. from ..fixer_util import Name, Call, Dot
  28. from .. import fixer_util
  29. iter_exempt = fixer_util.consuming_calls | {"iter"}
  30. class FixDict(fixer_base.BaseFix):
  31. BM_compatible = True
  32. PATTERN = """
  33. power< head=any+
  34. trailer< '.' method=('keys'|'items'|'values'|
  35. 'iterkeys'|'iteritems'|'itervalues'|
  36. 'viewkeys'|'viewitems'|'viewvalues') >
  37. parens=trailer< '(' ')' >
  38. tail=any*
  39. >
  40. """
  41. def transform(self, node, results):
  42. head = results["head"]
  43. method = results["method"][0] # Extract node for method name
  44. tail = results["tail"]
  45. syms = self.syms
  46. method_name = method.value
  47. isiter = method_name.startswith("iter")
  48. isview = method_name.startswith("view")
  49. if isiter or isview:
  50. method_name = method_name[4:]
  51. assert method_name in ("keys", "items", "values"), repr(method)
  52. head = [n.clone() for n in head]
  53. tail = [n.clone() for n in tail]
  54. special = not tail and self.in_special_context(node, isiter)
  55. args = head + [pytree.Node(syms.trailer,
  56. [Dot(),
  57. Name(method_name,
  58. prefix=method.prefix)]),
  59. results["parens"].clone()]
  60. new = pytree.Node(syms.power, args)
  61. if not (special or isview):
  62. new.prefix = ""
  63. new = Call(Name("iter" if isiter else "list"), [new])
  64. if tail:
  65. new = pytree.Node(syms.power, [new] + tail)
  66. new.prefix = node.prefix
  67. return new
  68. P1 = "power< func=NAME trailer< '(' node=any ')' > any* >"
  69. p1 = patcomp.compile_pattern(P1)
  70. P2 = """for_stmt< 'for' any 'in' node=any ':' any* >
  71. | comp_for< 'for' any 'in' node=any any* >
  72. """
  73. p2 = patcomp.compile_pattern(P2)
  74. def in_special_context(self, node, isiter):
  75. if node.parent is None:
  76. return False
  77. results = {}
  78. if (node.parent.parent is not None and
  79. self.p1.match(node.parent.parent, results) and
  80. results["node"] is node):
  81. if isiter:
  82. # iter(d.iterkeys()) -> iter(d.keys()), etc.
  83. return results["func"].value in iter_exempt
  84. else:
  85. # list(d.keys()) -> list(d.keys()), etc.
  86. return results["func"].value in fixer_util.consuming_calls
  87. if not isiter:
  88. return False
  89. # for ... in d.iterkeys() -> for ... in d.keys(), etc.
  90. return self.p2.match(node.parent, results) and results["node"] is node