fix_operator.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. """Fixer for operator functions.
  2. operator.isCallable(obj) -> callable(obj)
  3. operator.sequenceIncludes(obj) -> operator.contains(obj)
  4. operator.isSequenceType(obj) -> isinstance(obj, collections.abc.Sequence)
  5. operator.isMappingType(obj) -> isinstance(obj, collections.abc.Mapping)
  6. operator.isNumberType(obj) -> isinstance(obj, numbers.Number)
  7. operator.repeat(obj, n) -> operator.mul(obj, n)
  8. operator.irepeat(obj, n) -> operator.imul(obj, n)
  9. """
  10. import collections.abc
  11. # Local imports
  12. from lib2to3 import fixer_base
  13. from lib2to3.fixer_util import Call, Name, String, touch_import
  14. def invocation(s):
  15. def dec(f):
  16. f.invocation = s
  17. return f
  18. return dec
  19. class FixOperator(fixer_base.BaseFix):
  20. BM_compatible = True
  21. order = "pre"
  22. methods = """
  23. method=('isCallable'|'sequenceIncludes'
  24. |'isSequenceType'|'isMappingType'|'isNumberType'
  25. |'repeat'|'irepeat')
  26. """
  27. obj = "'(' obj=any ')'"
  28. PATTERN = """
  29. power< module='operator'
  30. trailer< '.' %(methods)s > trailer< %(obj)s > >
  31. |
  32. power< %(methods)s trailer< %(obj)s > >
  33. """ % dict(methods=methods, obj=obj)
  34. def transform(self, node, results):
  35. method = self._check_method(node, results)
  36. if method is not None:
  37. return method(node, results)
  38. @invocation("operator.contains(%s)")
  39. def _sequenceIncludes(self, node, results):
  40. return self._handle_rename(node, results, "contains")
  41. @invocation("callable(%s)")
  42. def _isCallable(self, node, results):
  43. obj = results["obj"]
  44. return Call(Name("callable"), [obj.clone()], prefix=node.prefix)
  45. @invocation("operator.mul(%s)")
  46. def _repeat(self, node, results):
  47. return self._handle_rename(node, results, "mul")
  48. @invocation("operator.imul(%s)")
  49. def _irepeat(self, node, results):
  50. return self._handle_rename(node, results, "imul")
  51. @invocation("isinstance(%s, collections.abc.Sequence)")
  52. def _isSequenceType(self, node, results):
  53. return self._handle_type2abc(node, results, "collections.abc", "Sequence")
  54. @invocation("isinstance(%s, collections.abc.Mapping)")
  55. def _isMappingType(self, node, results):
  56. return self._handle_type2abc(node, results, "collections.abc", "Mapping")
  57. @invocation("isinstance(%s, numbers.Number)")
  58. def _isNumberType(self, node, results):
  59. return self._handle_type2abc(node, results, "numbers", "Number")
  60. def _handle_rename(self, node, results, name):
  61. method = results["method"][0]
  62. method.value = name
  63. method.changed()
  64. def _handle_type2abc(self, node, results, module, abc):
  65. touch_import(None, module, node)
  66. obj = results["obj"]
  67. args = [obj.clone(), String(", " + ".".join([module, abc]))]
  68. return Call(Name("isinstance"), args, prefix=node.prefix)
  69. def _check_method(self, node, results):
  70. method = getattr(self, "_" + results["method"][0].value)
  71. if isinstance(method, collections.abc.Callable):
  72. if "module" in results:
  73. return method
  74. else:
  75. sub = (str(results["obj"]),)
  76. invocation_str = method.invocation % sub
  77. self.warning(node, "You should use '%s' here." % invocation_str)
  78. return None