import ast
import random

from .operator import Operator


class RemoveDecorator(Operator):
    """An operator that removes a single random decorator."""
    regular_decorators = ["classmethod", "staticmethod", "abstractmethod"]

    def visit_FunctionDef(self, node):  # noqa
        if [x for x in node.decorator_list if x.id not in self.regular_decorators]:
            return self.visit_mutation_site(node)
        return node

    def mutate(self, node, _):
        """Modify the decorator list to include randomly one less"""
        decorator_candidates = [x for x in node.decorator_list if x.id not in self.regular_decorators]
        removed_decorator = random.choice(decorator_candidates)
        modified_decorator = [x for x in node.decorator_list if x.id != removed_decorator.id]
        node.decorator_list = modified_decorator
        
        return node

import ast

from .operator import Operator


class RemoveDecorator(Operator):
    """An operator that removes each of the non standard decorators."""
    regular_decorators = ["classmethod", "staticmethod", "abstractmethod"]

    def visit_FunctionDef(self, node):  # noqa
        decorator_candidates = [x for x in node.decorator_list if \
                                x.id not in self.regular_decorators]
        if decorator_candidates:
            return self.visit_mutation_site(node, len(decorator_candidates))

        return node

    def mutate(self, node, idx):
        """Modify the decorator list to remove one decorator at each mutation"""
        decorator_candidates = [x for x in node.decorator_list if \
                                x.id not in self.regular_decorators]
        removed_decorator = decorator_candidates[idx]
        modified_decorator = [x for x in node.decorator_list if \
                              x.id != removed_decorator.id]
        node.decorator_list = modified_decorator
        
        return node

update remove decorator to remove each decorator iteratively