forked from philipperemy/name-dataset
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathquery.py
52 lines (40 loc) · 1.65 KB
/
query.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
import os
def _query(search_set, key, use_upper_case):
if use_upper_case and key.title() != key:
return False
return key.strip().lower() in search_set
class NameDataset:
FIRST_NAME_SEARCH = 0
LAST_NAME_SEARCH = 1
def __init__(self):
first_names_filename = os.path.join(os.path.dirname(__file__), 'first_names.all.txt')
with open(first_names_filename, 'rb') as r:
self.first_names = set(r.read().strip().split('\n'))
last_names_filename = os.path.join(os.path.dirname(__file__), 'last_names.all.txt')
with open(last_names_filename, 'rb') as r:
self.last_names = set(r.read().strip().split('\n'))
def search_first_name(self, first_name, use_upper_case=False):
return _query(self.first_names, first_name, use_upper_case)
def search_last_name(self, last_name, use_upper_case=False):
return _query(self.last_names, last_name, use_upper_case)
if __name__ == '__main__':
import sys
if sys.version_info < (3, 0):
print('Please use Python3+')
exit(1)
if len(sys.argv) < 2:
print('Give names separated by a comma as input.')
sys.exit(1)
m = NameDataset()
names_list = sys.argv[1].split(',')
print('----- First names ----')
print('Name'.ljust(30), 'Present?')
for name in names_list:
# ljust just for aesthetic reasons ;)
print(str(name).ljust(30), m.search_first_name(name))
print('----- Last names ----')
print('Name'.ljust(30), 'Present?')
for name in names_list:
# ljust just for aesthetic reasons ;)
print(str(name).ljust(30), m.search_last_name(name))
sys.exit(0)