GraphQL generic list using interfaces

test-graphql-generic-list.py

import graphene
import json
import sys


class Item(graphene.Interface):
    id = graphene.String(required=True)


class PaginatedList(graphene.Interface):
    items = graphene.List(Item, required=True)
    total_count = graphene.Int(required=True)


class Agent(graphene.ObjectType):
    class Meta:
        interfaces = (Item, )

    status = graphene.String()


class AgentList(graphene.ObjectType):
    class Meta:
        interfaces = (PaginatedList, )

    items = graphene.List(Agent, required=True)


agents_data = [
    Agent(id='test01', status='RUNNING'),
    Agent(id='test02', status='RUNNING'),
    Agent(id='test03', status='ERROR'),
    Agent(id='test04', status='STOPPED'),
    Agent(id='test05', status='RUNNING'),
    Agent(id='test06', status='STOPPED'),
]


class Query(graphene.ObjectType):
    agents = graphene.Field(
        AgentList,
        limit=graphene.Int(),
        offset=graphene.Int(),
        status=graphene.String())

    def resolve_agents(self, info, limit, offset, status=None):
        if status is None:
            count = len(agents_data)
            items = agents_data[offset:offset + limit]
        else:
            filtered_items = [
                item for item in agents_data
                if item.status == status]
            count = len(filtered_items)
            items = filtered_items[offset:offset + limit]
        return AgentList(items, total_count=count)


def do(offset, limit, status):
    schema = graphene.Schema(query=Query, auto_camelcase=False)
    result = schema.execute(
        'query($limit: Int!, $offset: Int!, $status: String) {\n'
        '  agents(limit: $limit, offset: $offset, status: $status) {\n'
        '    items { id status }\n'
        '    total_count\n'
        '  }\n'
        '}',
        variables={'offset': offset, 'limit': limit, 'status': status})
    if result.errors:
        print(result.errors)
    else:
        print(json.dumps(result.data, indent=2))


if __name__ == '__main__':
    do(int(sys.argv[1]), int(sys.argv[2]),
           sys.argv[3] if len(sys.argv) >= 4 else None)