Require data_source when creating a RelayMixin.

Author: jhgg

epoxy/contrib/relay/data_source/base.py

@@ -0,0 +1,7 @@
+class BaseDataSource(object):
+    def fetch_node(self, object_type, id, resolve_info):
+        raise NotImplementedError('Must implement fetch_node to resolve node by ID.')
+
+    def make_connection_resolver(self, relay, object_type_thunk):
+        raise NotImplementedError('Must implement make_connection_resolver so that RelayMixin can automatically '
+                                  'create connection resolvers')

epoxy/contrib/relay/data_source/memory.py

@@ -1,10 +1,11 @@
 from collections import defaultdict
-
 from operator import attrgetter
+
 from ..connections.sorted_collection import SortedCollection
+from .base import BaseDataSource
 
 
-class InMemoryDataSource(object):
+class InMemoryDataSource(BaseDataSource):
     def __init__(self):
         self.objects_by_type_and_id = defaultdict(dict)
         self.objects_by_type = defaultdict(lambda: SortedCollection(key=attrgetter('id')))
@@ -16,9 +17,6 @@ def add(self, obj):
     def fetch_node(self, object_type, id, resolve_info):
         return self.objects_by_type_and_id[object_type].get(id)
 
-    def get_collection(self, object_type):
-        return self.objects_by_type[object_type()]
-
     def make_connection_resolver(self, relay, object_type_thunk):
         def resolver(obj, args, info):
             object_type = relay.R[object_type_thunk]()

epoxy/contrib/relay/mixin.py

@@ -5,11 +5,11 @@
 
 
 class RelayMixin(object):
-    def __init__(self, registry):
+    def __init__(self, registry, data_source):
         self.R = registry
+        self.data_source = data_source
         self._node_field = None
         self._connections = {}
-        self.data_source = None
 
     @property
     def NodeField(self):
@@ -25,9 +25,6 @@ def NodeField(self):
             resolver=lambda obj, args, info: self.fetch_node(args.get('id'), info)
         )
 
-    def set_data_source(self, data_source):
-        self.data_source = data_source
-
     def get_connection_and_edge_types(self, type_name):
         return self._connections[type_name]
 
@@ -78,8 +75,11 @@ class Connection(R.ObjectType):
         self._connections[name] = Connection, Edge
         return Connection, Edge
 
-    def Connection(self, name, object_type, args=None, **kwargs):
+    def Connection(self, name, object_type, args=None, resolver=None, **kwargs):
         args = args or {}
         args.update(connection_args)
-        field = self.R.Field(self.connection_definitions(name, object_type)[0], args=args, **kwargs)
+        if not resolver:
+            resolver = self.data_source.make_connection_resolver(self, object_type)
+
+        field = self.R.Field(self.connection_definitions(name, object_type)[0], args=args, resolver=resolver, **kwargs)
         return field

epoxy/registry.py

@@ -219,8 +219,8 @@ def Schema(self, query, mutation=None):
         self._add_impl_to_interfaces()
         return GraphQLSchema(query=query, mutation=mutation)
 
-    def Mixin(self, mixin_cls):
-        mixin = mixin_cls(self)
+    def Mixin(self, mixin_cls, *args, **kwargs):
+        mixin = mixin_cls(self, *args, **kwargs)
         mixin.register_types()
         return mixin
 

tests/test_relay/test_relay_connections.py

@@ -14,19 +14,18 @@
              "Spike", "Gunner", "Sparky", "Mickey", "Kobe", "Chase", "Oreo", "Frankie", "Mac", "Benji", "Bubba",
              "Champ", "Brady", "Elvis", "Copper", "Cash", "Archie", "Walter"]
 
+data_source = InMemoryDataSource()
+
 R = TypeRegistry()
-Relay = R.Mixin(RelayMixin)
+Relay = R.Mixin(RelayMixin, data_source)
 
 
-class Pet(R.Implements[R.Node]):
+class Pet(R.Implements[Relay.Node]):
     name = R.String
 
 
-data_source = InMemoryDataSource()
-
-
 class Query(R.ObjectType):
-    pets = Relay.Connection('Pet', R.Pet, resolver=data_source.make_connection_resolver(Relay, R.Pet))
+    pets = Relay.Connection('Pet', R.Pet)
     node = Relay.NodeField
 
 

tests/test_relay/test_relay_node.py

@@ -6,7 +6,7 @@
 
 def test_relay_node_definition():
     R = TypeRegistry()
-    Relay = R.Mixin(RelayMixin)
+    Relay = R.Mixin(RelayMixin, InMemoryDataSource())
 
     class Pet(R.Implements[R.Node]):
         name = R.String
@@ -30,7 +30,7 @@ class Query(R.ObjectType):
 
 def test_relay_node_definition_using_custom_type():
     R = TypeRegistry()
-    Relay = R.Mixin(RelayMixin)
+    Relay = R.Mixin(RelayMixin, InMemoryDataSource())
 
     class Pet(R.Implements[R.Node]):
         name = R.String
@@ -59,10 +59,10 @@ def __init__(self, id, name):
 
 
 def test_relay_node_field_resolver():
-    R = TypeRegistry()
-    Relay = R.Mixin(RelayMixin)
     data_source = InMemoryDataSource()
-    Relay.set_data_source(data_source)
+
+    R = TypeRegistry()
+    Relay = R.Mixin(RelayMixin, data_source)
 
     class Pet(R.Implements[R.Node]):
         name = R.String