From cd6194e70d4cd2c778dc65a9c1d5e10d9899eda6 Mon Sep 17 00:00:00 2001 From: Andrew Bates Date: Tue, 21 May 2024 10:51:12 -0400 Subject: [PATCH] feat: Now allowing direct assignment for many-to-many or using through --- .../testdata/next_prefix_for_location.yaml | 42 ++++++++++ nautobot_design_builder/fields.py | 78 +++++++++++++++---- .../tests/testdata/prefixes_for_location.yaml | 23 ++++++ ...prefixes_for_location_without_through.yaml | 21 +++++ 4 files changed, 147 insertions(+), 17 deletions(-) create mode 100644 nautobot_design_builder/contrib/tests/testdata/next_prefix_for_location.yaml create mode 100644 nautobot_design_builder/tests/testdata/prefixes_for_location.yaml create mode 100644 nautobot_design_builder/tests/testdata/prefixes_for_location_without_through.yaml diff --git a/nautobot_design_builder/contrib/tests/testdata/next_prefix_for_location.yaml b/nautobot_design_builder/contrib/tests/testdata/next_prefix_for_location.yaml new file mode 100644 index 00000000..8ba4d086 --- /dev/null +++ b/nautobot_design_builder/contrib/tests/testdata/next_prefix_for_location.yaml @@ -0,0 +1,42 @@ +--- +extensions: + - "nautobot_design_builder.contrib.ext.NextPrefixExtension" + - "nautobot_design_builder.contrib.ext.ChildPrefixExtension" +designs: + - location_types: + - name: "Region" + content_types: + - "!get:app_label": "ipam" + "!get:model": "prefix" + locations: + - "name": "Region" + "location_type__name": "Region" + "status__name": "Active" + prefixes: + - prefix: "10.0.0.0/23" + type: "container" + status__name: "Active" + locations: + - location: + "!get:name": "Region" + - "!next_prefix": + locations__name: "Region" + type: "container" + length: 26 + status__name: "Active" + description: "Region Parent Prefix" + type: "container" + - "!next_prefix": + locations__name: "Region" + type: "container" + length: 26 + status__name: "Active" + description: "Region Parent Prefix" + type: "container" +checks: + - model_exists: + model: "nautobot.ipam.models.Prefix" + query: {prefix: "10.0.0.0/26"} + - model_exists: + model: "nautobot.ipam.models.Prefix" + query: {prefix: "10.0.0.64/26"} diff --git a/nautobot_design_builder/fields.py b/nautobot_design_builder/fields.py index c9b59713..77d6bd8c 100644 --- a/nautobot_design_builder/fields.py +++ b/nautobot_design_builder/fields.py @@ -145,7 +145,9 @@ class RelationshipFieldMixin: # pylint:disable=too-few-public-methods Relationship fields also include the reverse side of fields or even custom relationships. """ - def _get_instance(self, obj: "ModelInstance", value: Any, relationship_manager: "Manager" = None): + def _get_instance( + self, obj: "ModelInstance", value: Any, relationship_manager: "Manager" = None, related_model=None + ): """Helper function to create a new child model from a value. If the passed-in value is a dictionary, this method assumes that the dictionary @@ -165,8 +167,10 @@ def _get_instance(self, obj: "ModelInstance", value: Any, relationship_manager: Returns: ModelInstance: Either a newly created `ModelInstance` or the original value. """ + if related_model is None: + related_model = self.related_model if isinstance(value, Mapping): - value = obj.create_child(self.related_model, value, relationship_manager) + value = obj.create_child(related_model, value, relationship_manager) return value @@ -214,28 +218,60 @@ class ManyToManyField(BaseModelField, RelationshipFieldMixin): # pylint:disable def __init__(self, field: django_models.Field): # noqa:D102 super().__init__(field) self.auto_through = True - self.through_fields = field.remote_field.through_fields - through = field.remote_field.through - if not through._meta.auto_created: + self._init_through() + + def _init_through(self): + self.through = self.field.remote_field.through + if not self.through._meta.auto_created: self.auto_through = False - self.related_model = through - if field.remote_field.through_fields: - self.link_field = field.remote_field.through_fields[0] + if self.field.remote_field.through_fields: + self.link_field = self.field.remote_field.through_fields[0] else: - for f in through._meta.fields: - if f.related_model == field.model: + for f in self.through._meta.fields: + if f.related_model == self.field.model: self.link_field = f.name + def _get_related_model(self, value): + """Get the appropriate related model for the value. + + if there is an explicit through class, then we have two choices: + 1) Assign explicitly using the through-class attributes + 2) Assign implicitly like a normal many-to-many + + We want to be able to handle both situations, because it may be that + the through class has additional attributes. The way we determine if + the design is requesting the through-class or the implicit related class + is by examining the values to be assigned and matching their keys with + the related model and through model. + """ + if isinstance(value, Mapping): + attributes = set() + # Extract all of the top-level field names from the query in order + # to match them against available fields in the through table. If + # the set of attributes is a subset of the through class's attributes + # then use the through class directly, otherwise use the related_model + # class + for attribute in value.keys(): + if attribute.startswith("!get") or attribute.startswith("!create"): + attribute_parts = attribute.split(":") + attribute = attribute_parts[1] + + if "__" in attribute: + attribute = attribute.split("__")[0] + attributes.add(attribute) + through_fields = set(field.name for field in self.through._meta.fields) + if self.auto_through is False and attributes.issubset(through_fields): + return self.through + return self.related_model + @debug_set def __set__(self, obj: "ModelInstance", values): # noqa:D105 def setter(): items = [] for value in values: - value = self._get_instance(obj, value, getattr(obj.instance, self.field_name)) - if self.auto_through: - # Only need to call `add` if the through relationship was - # auto-created. Otherwise we explicitly create the through - # object + related_model = self._get_related_model(value) + value = self._get_instance(obj, value, getattr(obj.instance, self.field_name), related_model) + if related_model is not self.through: items.append(value.instance) else: setattr(value.instance, self.link_field, obj.instance) @@ -252,8 +288,16 @@ def setter(): class ManyToManyRelField(ManyToManyField): # pylint:disable=too-few-public-methods """Reverse many to many relationship field.""" - def __init__(self, field: django_models.Field): # noqa:D102 - super().__init__(field.remote_field) + def _init_through(self): + self.through = self.field.through + if not self.through._meta.auto_created: + self.auto_through = False + if self.field.through_fields: + self.link_field = self.field.through_fields[0] + else: + for f in self.through._meta.fields: + if f.related_model == self.field.model: + self.link_field = f.name class GenericRelationField(BaseModelField, RelationshipFieldMixin): # pylint:disable=too-few-public-methods diff --git a/nautobot_design_builder/tests/testdata/prefixes_for_location.yaml b/nautobot_design_builder/tests/testdata/prefixes_for_location.yaml new file mode 100644 index 00000000..a862c911 --- /dev/null +++ b/nautobot_design_builder/tests/testdata/prefixes_for_location.yaml @@ -0,0 +1,23 @@ +--- +depends_on: "base_test.yaml" +designs: + - locations: + - "!update:name": "Site" + prefixes: + - prefix: + prefix: "10.1.0.0/16" + status__name: "Active" + - prefixes: + - prefix: "10.0.0.0/23" + type: "container" + status__name: "Active" + locations: + - location: + "!get:name": "Site" +checks: + - model_exists: + model: "nautobot.ipam.models.Prefix" + query: {location__name: "Site", prefix: "10.0.0.0/23"} + - model_exists: + model: "nautobot.ipam.models.Prefix" + query: {location__name: "Site", prefix: "10.1.0.0/16"} diff --git a/nautobot_design_builder/tests/testdata/prefixes_for_location_without_through.yaml b/nautobot_design_builder/tests/testdata/prefixes_for_location_without_through.yaml new file mode 100644 index 00000000..2dae7b84 --- /dev/null +++ b/nautobot_design_builder/tests/testdata/prefixes_for_location_without_through.yaml @@ -0,0 +1,21 @@ +--- +depends_on: "base_test.yaml" +designs: + - locations: + - "!update:name": "Site" + prefixes: + - prefix: "10.1.0.0/16" + status__name: "Active" + - prefixes: + - prefix: "10.0.0.0/23" + type: "container" + status__name: "Active" + locations: + - "!get:name": "Site" +checks: + - model_exists: + model: "nautobot.ipam.models.Prefix" + query: {location__name: "Site", prefix: "10.0.0.0/23"} + - model_exists: + model: "nautobot.ipam.models.Prefix" + query: {location__name: "Site", prefix: "10.1.0.0/16"}