from stagecraft.apps.datasets.models import DataGroup, DataSet, DataType


class DataSetMassUpdate():
    @classmethod
    def update_bearer_token_for_data_type_or_group_name(cls, query, new_token):
        model_filter = DataSet.objects
        if 'data_type' in query:
            data_type = cls._get_model_instance_by_name(
                DataType, query['data_type'])
            model_filter = model_filter.filter(data_type=data_type)
        if 'data_group' in query:
            data_group = cls._get_model_instance_by_name(
                DataGroup, query['data_group'])
            model_filter = model_filter.filter(data_group=data_group)

        model_filter.update(bearer_token=new_token)

    @classmethod
    def _get_model_instance_by_name(cls, model, name):
        return model.objects.get(name=name)

from stagecraft.apps.datasets.models import DataGroup, DataSet, DataType


class DataSetMassUpdate(object):

    @classmethod
    def update_bearer_token_for_data_type_or_group_name(cls, query, new_token):
        cls(query).update(bearer_token=new_token)

    def __init__(self, query_dict):
        self.model_filter = DataSet.objects

        if 'data_type' in query_dict:
            data_type = self._get_model_instance_by_name(
                DataType, query_dict['data_type'])
            self.model_filter = self.model_filter.filter(data_type=data_type)
        if 'data_group' in query_dict:
            data_group = self._get_model_instance_by_name(
                DataGroup, query_dict['data_group'])
            self.model_filter = self.model_filter.filter(data_group=data_group)

    def update(self, **kwargs):
        self.model_filter.update(**kwargs)

    def _get_model_instance_by_name(self, model, name):
        return model.objects.get(name=name)

Refactor query out into instance with delegates the update