Improved handling of decorator nodes without an 'id' field.
-import ast
-
from .operator import Operator
@@ -8,9 +6,15 @@ class RemoveDecorator(Operator):
REGULAR_DECORATORS = frozenset(["classmethod", "staticmethod",
"abstractmethod"])
+ @classmethod
+ def _skip_decorator(cls, node):
+ """Determine if a decorator node should not be mutated.
+ """
+ return hasattr(node, 'id') and node.id in cls.REGULAR_DECORATORS
+
def visit_FunctionDef(self, node): # noqa
decorator_candidates = [x for x in node.decorator_list
- if hasattr(x, 'id') and x.id not in self.REGULAR_DECORATORS]
+ if not self._skip_decorator(x)]
if decorator_candidates:
return self.visit_mutation_site(node, len(decorator_candidates))
@@ -19,7 +23,7 @@ class RemoveDecorator(Operator):
def mutate(self, node, idx):
"""Modify the decorator list to remove one decorator at each mutation"""
candidates = [(i, d) for (i, d) in enumerate(node.decorator_list)
- if d.id not in self.REGULAR_DECORATORS]
+ if not self._skip_decorator(d)]
remove_idx, _ = candidates[idx]
del node.decorator_list[remove_idx]
return node
