Used Graphene types for relay Connection and Edges

Author: syrusakbary

graphene/core/types.py

@@ -30,7 +30,7 @@ def __new__(cls, name, bases, attrs):
             # If this isn't a subclass of Model, don't do anything special.
             return super_new(cls, name, bases, attrs)
 
-        module = attrs.pop('__module__')
+        module = attrs.pop('__module__', None)
         doc = attrs.pop('__doc__', None)
         new_class = super_new(cls, name, bases, {
             '__module__': module,

graphene/relay/__init__.py

@@ -4,6 +4,11 @@
     GlobalIDField,
 )
 
-from graphene.relay.types import Node
+from graphene.relay.types import (
+	Node,
+	PageInfo,
+	Edge,
+	Connection
+)
 
 from graphene.relay.utils import is_node

graphene/relay/fields.py

@@ -21,9 +21,14 @@
 
 class ConnectionField(Field):
 
-    def __init__(self, field_type, resolve=None, description=''):
+    def __init__(self, field_type, resolve=None, description='', connection_type=None, edge_type=None, **kwargs):
+        from graphene.relay.types import Connection, Edge
         super(ConnectionField, self).__init__(field_type, resolve=resolve,
-                                              args=connectionArgs, description=description)
+                                              args=connectionArgs, description=description, **kwargs)
+        self.connection_type = connection_type or Connection
+        self.edge_type = edge_type or Edge
+        assert issubclass(self.connection_type, Connection), 'connection_type in %r must be a subclass of Connection' % self
+        assert issubclass(self.edge_type, Edge), 'edge_type in %r must be a subclass of Edge' % self
 
     def wrap_resolved(self, value, instance, args, info):
         return value
@@ -39,9 +44,12 @@ def resolve(self, instance, args, info):
     @memoize
     def internal_type(self, schema):
         from graphene.relay.utils import is_node
-        object_type = self.get_object_type(schema)
-        assert is_node(object_type), 'Only nodes have connections.'
-        return object_type.get_connection(schema)
+        node = self.get_object_type(schema)
+        assert is_node(node), 'Only nodes have connections.'
+        schema.register(node)
+        edge_node_type = self.edge_type.for_node(node)
+        connection_node_type = self.connection_type.for_node(node, edge_type=edge_node_type)
+        return connection_node_type.internal_type(schema)
 
 
 class NodeField(Field):

graphene/relay/types.py

@@ -5,7 +5,8 @@
     connection_definitions
 )
 
-from graphene.core.types import Interface
+from graphene.core.types import Interface, ObjectType
+from graphene.core.fields import BooleanField, StringField, ListField, Field
 from graphene.relay.fields import GlobalIDField
 from graphene.utils import memoize
 
@@ -35,3 +36,37 @@ def to_global_id(cls, instance, args, info):
 class Node(BaseNode, Interface):
     '''An object with an ID'''
     id = GlobalIDField()
+
+
+class PageInfo(ObjectType):
+    has_next_page = BooleanField(required=True, description='When paginating forwards, are there more items?')
+    has_previous_page = BooleanField(required=True, description='When paginating backwards, are there more items?')
+    start_cursor = StringField(description='When paginating backwards, the cursor to continue.')
+    end_cursor = StringField(description='When paginating forwards, the cursor to continue.')
+
+
+class Edge(ObjectType):
+    '''An edge in a connection.'''
+    node = Field(lambda field: field.object_type.node_type, description='The item at the end of the edge')
+    end_cursor = StringField(required=True, description='A cursor for use in pagination')
+
+    @classmethod
+    @memoize
+    def for_node(cls, node):
+        from graphene.relay.utils import is_node
+        assert is_node(node), 'ObjectTypes in a edge have to be Nodes'
+        return type('%sEdge' % node._meta.type_name, (cls, ), {'node_type': node})
+
+
+class Connection(ObjectType):
+    '''A connection to a list of items.'''
+    page_info = Field(PageInfo, required=True, description='The Information to aid in pagination')
+    edges = ListField(lambda field: field.object_type.edge_type, description='Information to aid in pagination.')
+
+    @classmethod
+    @memoize
+    def for_node(cls, node, edge_type=None):
+        from graphene.relay.utils import is_node
+        edge_type = edge_type or Edge
+        assert is_node(node), 'ObjectTypes in a connection have to be Nodes'
+        return type('%sConnection' % node._meta.type_name, (cls, ), {'edge_type': edge_type.for_node(node)})

tests/relay/test_relayfields.py

@@ -12,6 +12,11 @@
 
 class MyType(object):
     name = 'my'
+    arg = None
+
+
+class MyConnection(relay.Connection):
+    my_custom_field = graphene.StringField(resolve=lambda instance, *_: 'Custom')
 
 
 class MyNode(relay.Node):
@@ -24,7 +29,13 @@ def get_node(cls, id):
 
 class Query(graphene.ObjectType):
     my_node = relay.NodeField(MyNode)
+    all_my_nodes = relay.ConnectionField(MyNode, connection_type=MyConnection, customArg=graphene.Argument(graphene.String))
 
+    def resolve_all_my_nodes(self, args, info):
+        t = MyType()
+        custom_arg = args.get('customArg')
+        assert custom_arg == "1"
+        return [MyNode(t)]
 
 schema.query = Query
 
@@ -34,12 +45,34 @@ def test_nodefield_query():
     query RebelsShipsQuery {
       myNode(id:"TXlOb2RlOjE=") {
         name
+      },
+      allMyNodes (customArg:"1") {
+        edges {
+          node {
+            name
+          }
+        },
+        myCustomField
+        pageInfo {
+          hasNextPage
+        }
       }
     }
     '''
     expected = {
         'myNode': {
             'name': 'my'
+        },
+        'allMyNodes': {
+            'edges': [{
+                'node': {
+                    'name': 'my'
+                }
+            }],
+            'myCustomField': 'Custom',
+            'pageInfo': {
+                'hasNextPage': False,
+            }
         }
     }
     result = schema.execute(query)