graphene/examples/starwars/data.py

human_data = {}
droid_data = {}


def setup():
    from .schema import Human, Droid
    global human_data, droid_data
    luke = Human(
        id='1000',
        name='Luke Skywalker',
        friends=['1002', '1003', '2000', '2001'],
        appears_in=[4, 5, 6],
        home_planet='Tatooine',
    )

    vader = Human(
        id='1001',
        name='Darth Vader',
        friends=['1004'],
        appears_in=[4, 5, 6],
        home_planet='Tatooine',
    )

    han = Human(
        id='1002',
        name='Han Solo',
        friends=['1000', '1003', '2001'],
        appears_in=[4, 5, 6],
        home_planet=None,
    )

    leia = Human(
        id='1003',
        name='Leia Organa',
        friends=['1000', '1002', '2000', '2001'],
        appears_in=[4, 5, 6],
        home_planet='Alderaan',
    )

    tarkin = Human(
        id='1004',
        name='Wilhuff Tarkin',
        friends=['1001'],
        appears_in=[4],
        home_planet=None,
    )

    human_data = {
        '1000': luke,
        '1001': vader,
        '1002': han,
        '1003': leia,
        '1004': tarkin,
    }

    c3po = Droid(
        id='2000',
        name='C-3PO',
        friends=['1000', '1002', '1003', '2001'],
        appears_in=[4, 5, 6],
        primary_function='Protocol',
    )

    r2d2 = Droid(
        id='2001',
        name='R2-D2',
        friends=['1000', '1002', '1003'],
        appears_in=[4, 5, 6],
        primary_function='Astromech',
    )

    droid_data = {
        '2000': c3po,
        '2001': r2d2,
    }


def get_character(id):
    return human_data.get(id) or droid_data.get(id)


def get_friends(character):
    return map(get_character, character.friends)


def get_hero(episode):
    if episode == 5:
        return human_data['1000']
    return droid_data['2001']


def get_human(id):
    return human_data.get(id)


def get_droid(id):
    return droid_data.get(id)