diff --git a/readthedocs/integrations/models.py b/readthedocs/integrations/models.py index 0830c4f71e3..785f4fba728 100644 --- a/readthedocs/integrations/models.py +++ b/readthedocs/integrations/models.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - """Integration models for external services.""" import json @@ -169,34 +167,42 @@ def formatted_response_body(self): return self.formatted_json('response_body') -class IntegrationQuerySet(models.QuerySet): +class SubclassMixin: """ - Return a subclass of Integration, based on the integration type. + Mixin used to return a subclass of a Model, based on `subclass_field`. + + `subclass_field_id` is the name of the field used as the identifier + for each subclass. It should be unique for each subclass. + And should be an value existing in `subclass_field`. .. note:: This doesn't affect queries currently, only fetching of an object """ - def _get_subclass(self, integration_type): - # Build a mapping of integration_type -> class dynamically + subclass_field = None + subclass_field_id = None + + def _get_subclass(self, subclass_id): + """Build a mapping of subclass_field_id -> class dynamically.""" class_map = { - cls.integration_type_id: cls + getattr(cls, self.subclass_field_id): cls for cls in self.model.__subclasses__() - if hasattr(cls, 'integration_type_id') - } # yapf: disable - return class_map.get(integration_type) + if hasattr(cls, self.subclass_field_id) + } + return class_map.get(subclass_id) def _get_subclass_replacement(self, original): """ - Replace model instance on Integration subclasses. + Replace model instance on a Model subclasses. - This is based on the ``integration_type`` field, and is used to provide - specific functionality to and integration via a proxy subclass of the - Integration model. + This is based on the ``subclass_field`` field, and is used to provide + specific functionality to the models via a proxy subclass. """ - cls_replace = self._get_subclass(original.integration_type) + cls_replace = self._get_subclass( + getattr(original, self.subclass_field) + ) if cls_replace is None: return original new = cls_replace() @@ -215,11 +221,11 @@ def create(self, **kwargs): """ Override of create method to use subclass instance instead. - Instead of using the underlying Integration model to create this - instance, we get the correct subclass to use instead. This allows for + Instead of using the underlying model to create this instance, + we get the correct subclass to use instead. This allows for overrides to ``save`` and other model functions on object creation. """ - model_cls = self._get_subclass(kwargs.get('integration_type')) + model_cls = self._get_subclass(kwargs.get(self.subclass_field)) if model_cls is None: model_cls = self.model obj = model_cls(**kwargs) @@ -228,6 +234,14 @@ def create(self, **kwargs): return obj +class IntegrationQuerySet(SubclassMixin, models.QuerySet): + + """Return a subclass of Integration, based on the integration type.""" + + subclass_field = 'integration_type' + subclass_field_id = 'integration_type_id' + + class Integration(models.Model): """Inbound webhook integration for projects."""