diff --git a/nautobot_ssot/contrib.py b/nautobot_ssot/contrib.py index d9abf26dd..1ad9dc3a9 100644 --- a/nautobot_ssot/contrib.py +++ b/nautobot_ssot/contrib.py @@ -249,9 +249,8 @@ def _handle_custom_relationship_to_many_relationship( ): # Introspect type annotations to deduce which fields are of interest # for this many-to-many relationship. - diffsync_field_type = diffsync_model.__annotations__[parameter_name] - # TODO: Why is this different then in the normal case?? - inner_type = diffsync_field_type.__dict__["__args__"][0].__dict__["__args__"][0] + diffsync_field_type = get_type_hints(diffsync_model)[parameter_name] + inner_type = diffsync_field_type.__dict__["__args__"][0] related_objects_list = [] # TODO: Allow for filtering, i.e. not taking into account all the objects behind the relationship. relationship = self.get_from_orm_cache({"name": annotation.name}, Relationship) @@ -280,7 +279,7 @@ def _handle_custom_relationship_to_many_relationship( association, "source" if annotation.side == RelationshipSideEnum.DESTINATION else "destination" ) dictionary_representation = { - field_name: getattr(related_object, field_name) for field_name in inner_type.__annotations__ + field_name: getattr(related_object, field_name) for field_name in get_type_hints(inner_type) } # Only use those where there is a single field defined, all 'None's will not help us. if any(dictionary_representation.values()): @@ -352,13 +351,13 @@ class NautobotInterface(NautobotModel): """ # Introspect type annotations to deduce which fields are of interest # for this many-to-many relationship. - diffsync_field_type = diffsync_model.__annotations__[parameter_name] + diffsync_field_type = get_type_hints(diffsync_model)[parameter_name] inner_type = diffsync_field_type.__dict__["__args__"][0] related_objects_list = [] # TODO: Allow for filtering, i.e. not taking into account all the objects behind the relationship. for related_object in getattr(database_object, parameter_name).all(): dictionary_representation = { - field_name: getattr(related_object, field_name) for field_name in inner_type.__annotations__ + field_name: getattr(related_object, field_name) for field_name in get_type_hints(inner_type) } # Only use those where there is a single field defined, all 'None's will not help us. if any(dictionary_representation.values()): diff --git a/nautobot_ssot/tests/test_contrib.py b/nautobot_ssot/tests/test_contrib.py index 2e5a9b159..0bd3d59d8 100644 --- a/nautobot_ssot/tests/test_contrib.py +++ b/nautobot_ssot/tests/test_contrib.py @@ -887,3 +887,43 @@ class _ProviderTestModel(NautobotModel): diffsync_provider.diffsync = NautobotAdapter(job=None) self.assertEqual(self.provider, diffsync_provider.get_from_db()) + + +class AnnotationsSubclassingTest(TestCase): + """Test that annotations work properly with subclassing.""" + + def test_annotations_subclassing(self): + """Test that annotations work properly with subclassing.""" + + class BaseTenantModel(NautobotModel): + """Tenant model to be subclassed.""" + + _model = tenancy_models.Tenant + _modelname = "tenant" + _identifiers = ("name",) + _attributes = ("tags",) + + name: str + tags: List[TagDict] + + class Subclass(BaseTenantModel): + """Subclassed model.""" + + extra_field: Optional[str] = None + + class Adapter(NautobotAdapter): + """Test adapter.""" + + tenant = Subclass + top_level = ["tenant"] + + tenancy_models.Tenant.objects.create(name="Test Tenant") + + adapter = Adapter(job=None) + try: + adapter.load() + except KeyError as error: + if error.args[0] == "tags": + self.fail("Don't use `Klass.__annotations__`, prefer `typing.get_type_hints`.") + else: + raise error