-
Notifications
You must be signed in to change notification settings - Fork 51
/
species.py
57 lines (43 loc) · 1.8 KB
/
species.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
# Copyright 2018-present Kensho Technologies, LLC.
import random
from .utils import create_edge_statement, create_vertex_statement, get_random_limbs, get_uuid
SPECIES_LIST = (
'Nazgul',
'Pteranodon',
'Dragon',
'Hippogriff',
)
FOOD_LIST = (
"Bacon",
"Lembas",
"Bloood pie",
)
NUM_FOODS = 2
def _create_food_statement(food_name):
"""Return a SQL statement to create a Food vertex."""
field_name_to_value = {'name': food_name, 'uuid': get_uuid()}
return create_vertex_statement('Food', field_name_to_value)
def _create_species_statement(species_name):
"""Return a SQL statement to create a Species vertex."""
field_name_to_value = {'name': species_name, 'limbs': get_random_limbs(), 'uuid': get_uuid()}
return create_vertex_statement('Species', field_name_to_value)
def _create_species_eats_statement(from_name, to_name):
"""Return a SQL statement to create a Species_Eats edge."""
if to_name in SPECIES_LIST:
to_class = 'Species'
elif to_name in FOOD_LIST:
to_class = 'Food'
else:
raise AssertionError(u'Invalid name for Species_Eats endpoint: {}'.format(to_name))
return create_edge_statement('Species_Eats', 'Species', from_name, to_class, to_name)
def get_species_generation_commands():
"""Return a list of SQL statements to create all species vertices."""
command_list = []
for food_name in FOOD_LIST:
command_list.append(_create_food_statement(food_name))
for species_name in SPECIES_LIST:
command_list.append(_create_species_statement(species_name))
for species_name in SPECIES_LIST:
for food_or_species_name in random.sample(SPECIES_LIST + FOOD_LIST, NUM_FOODS):
command_list.append(_create_species_eats_statement(species_name, food_or_species_name))
return command_list