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