-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
135 lines (111 loc) · 3.12 KB
/
main.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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
from PyInquirer import prompt
from prompt_toolkit.validation import Validator, ValidationError
from pyfiglet import Figlet
from youtube.video_search import SearchQuery
from collector import collect_yt, collect_local
class NumberValidator(Validator):
def validate(self, document):
try:
int(document.text)
except ValueError:
raise ValidationError(
message='Please enter a number',
cursor_position=len(document.text))
def main():
'''
Print logo and begin questions with keyword.
'''
fig = Figlet(font='big')
print(fig.renderText('Keyword Collector'))
ask_keyword()
def ask_keyword():
'''
Ask user what keyword should be collected.
'''
questions = [
{
'type': 'input',
'name': 'keyword',
'message': 'What keyword would you like to collect?'
}
]
answers = prompt(questions)
keyword = answers['keyword']
user_input['keyword'] = keyword.strip()
ask_audio_source()
def ask_audio_source():
'''
Ask user where keywords should be found.
Options:
YouTube
Directory containing local WAV files
'''
questions = [
{
'type': 'list',
'name': 'audio_source',
'message': 'Where would you like to collect keywords from?',
'choices': [
{
'key': '1',
'name': 'YouTube',
'value': 'youtube'
},
{
'key': '2',
'name': 'Local WAV files',
'value': 'local'
}
]
}
]
answers = prompt(questions)
audio_source = answers['audio_source']
if audio_source == 'youtube':
ask_yt_options()
else:
ask_local_options()
def ask_yt_options():
'''
Ask user to enter YouTube options:
- Desired number of keywords
- Search query
'''
questions = [
{
'type': 'input',
'name': 'query',
'message': 'What would you like to search for on YouTube?'
},
{
'type': 'input',
'name': 'quantity',
'message': 'How many utterances would you like to collect?',
'validate': NumberValidator
}
]
answers = prompt(questions)
query = answers['query']
quantity = int(answers['quantity'])
keyword = user_input['keyword']
search_query = SearchQuery(keyword=keyword, query=query)
search_results = search_query.search()
collect_yt(search_results, keyword, quantity)
def ask_local_options():
'''
Ask user location of directory containing WAV files.
'''
questions = [
{
'type': 'input',
'name': 'data_dir',
'message': 'Enter path to directory containing WAV files:'
}
]
answers = prompt(questions)
data_dir = answers['data_dir']
keyword = user_input['keyword']
collect_local(data_dir, keyword)
if __name__ == "__main__":
user_input = {}
main()