Simplify overriding by adding a new method "maybe_optimize"

graphene_django_optimizer/types.py

 from graphene.types.definitions import GrapheneObjectType
 from graphene_django.types import DjangoObjectType
-from graphql import ResolveInfo
 
 from .query import query
 
@@ -10,20 +9,24 @@ class OptimizedDjangoObjectType(DjangoObjectType):
         abstract = True
 
     @classmethod
-    def can_optimize_resolver(cls, resolver_info: ResolveInfo):
+    def can_optimize_resolver(cls, resolver_info):
         return (
-            isinstance(resolver_info.return_type, GrapheneObjectType) and
-            resolver_info.return_type.graphene_type is cls)
+            isinstance(resolver_info.return_type, GrapheneObjectType)
+            and resolver_info.return_type.graphene_type is cls)
 
     @classmethod
     def get_optimized_node(cls, info, qs, pk):
+        return query(qs, info).get(pk=pk)
+
+    @classmethod
+    def maybe_optimize(cls, info, qs, pk):
         try:
-            return query(qs, info).get(pk=pk)
+            if cls.can_optimize_resolver(info):
+                return cls.get_optimized_node(info, qs, pk)
+            return qs.get(pk=pk)
         except cls._meta.model.DoesNotExist:
             return None
 
     @classmethod
-    def get_node(cls, info: ResolveInfo, id):
-        if cls.can_optimize_resolver(info):
-            return cls.get_optimized_node(info, cls._meta.model.objects, id)
-        return super(OptimizedDjangoObjectType, cls).get_node(info, id)
+    def get_node(cls, info, id):
+        return cls.maybe_optimize(info, cls._meta.model.objects, id)

tests/test_types.py

@@ -52,11 +52,11 @@ def test_should_return_none_when_node_is_not_resolved(mocked_optimizer):
     ''')
 
     info.return_type = schema.get_type('SomeOtherItemType')
-    qs = SomeOtherItem.objects.filter(name='foo')
+    qs = SomeOtherItem.objects
     mocked_optimizer.return_value = qs
 
-    assert SomeOtherItemType.get_optimized_node(info, qs, 7) is None
-    mocked_optimizer.assert_called_once_with(qs, info)
+    assert SomeOtherItemType.get_node(info, 8) is None
+    mocked_optimizer.assert_called_once_with(SomeOtherItem.objects, info)
 
 
 @patch('graphene_django_optimizer.types.query')