1
1
from cmd import Cmd
2
- from objectbox import *
2
+ from objectbox import Entity , Float32Vector , HnswIndex , Id , Store , String
3
3
import time
4
4
import csv
5
5
import os
6
6
7
+
7
8
@Entity ()
8
9
class City :
9
10
id = Id ()
10
11
name = String ()
11
- location = Float32Vector (index = HnswIndex (
12
- dimensions = 2 ,
13
- distance_type = VectorDistanceType .EUCLIDEAN
14
- ))
12
+ location = Float32Vector (index = HnswIndex (dimensions = 2 ))
13
+
15
14
16
15
def list_cities (cities ):
17
16
print ("{:3s} {:25s} {:>9s} {:>9s}" .format ("ID" , "Name" , "Latitude" , "Longitude" ))
18
17
for city in cities :
19
18
print ("{:3d} {:25s} {:>9.2f} {:>9.2f}" .format (
20
- city .id , city .name , city .location [0 ], city .location [1 ]))
19
+ city .id , city .name , city .location [0 ], city .location [1 ]))
20
+
21
21
22
22
def list_cities_with_scores (city_score_tuples ):
23
23
print ("{:3s} {:25s} {:>9s} {:>9s} {:>5s}" .format ("ID" , "Name" , "Latitude" , "Longitude" , "Score" ))
24
- for (city ,score ) in city_score_tuples :
24
+ for (city , score ) in city_score_tuples :
25
25
print ("{:3d} {:25s} {:>9.2f} {:>9.2f} {:>5.2f}" .format (
26
- city .id , city .name , city .location [0 ], city .location [1 ], score ))
26
+ city .id , city .name , city .location [0 ], city .location [1 ], score ))
27
+
27
28
28
29
class VectorSearchCitiesCmd (Cmd ):
29
30
prompt = "> "
31
+
30
32
def __init__ (self , * args ):
31
33
Cmd .__init__ (self , * args )
32
34
dbdir = "cities-db"
33
35
new_db = not os .path .exists (dbdir )
34
36
self ._store = Store (directory = dbdir )
35
37
self ._box = self ._store .box (City )
36
- if new_db :
38
+ if new_db :
37
39
with open (os .path .join (os .path .dirname (__file__ ), 'cities.csv' )) as f :
38
- r = csv .reader (f )
40
+ r = csv .reader (f )
39
41
cities = []
40
42
for row in r :
41
43
city = City ()
42
44
city .name = row [0 ]
43
- city .location = [ row [1 ], row [2 ] ]
45
+ city .location = [row [1 ], row [2 ]]
44
46
cities .append (city )
45
47
self ._box .put (* cities )
46
48
47
49
def do_ls (self , name : str = "" ):
48
50
"""list all cities or starting with <prefix>\n usage: ls [<prefix>]"""
49
- qb = self ._box .query ( City .name .starts_with (name ) )
51
+ qb = self ._box .query (City .name .starts_with (name ))
50
52
query = qb .build ()
51
53
list_cities (query .find ())
52
54
@@ -62,38 +64,38 @@ def do_city_neighbors(self, args: str):
62
64
num = 5
63
65
if len (args ) == 2 :
64
66
num = int (args [1 ])
65
- qb = self ._box .query ( City .name .equals (city ) )
67
+ qb = self ._box .query (City .name .equals (city ))
66
68
query = qb .build ()
67
69
cities = query .find ()
68
70
if len (cities ) == 1 :
69
71
location = cities [0 ].location
70
72
# +1 for the city
71
73
qb = self ._box .query (
72
- City .location .nearest_neighbor (location , num + 1 ) & City .name .not_equals (city )
74
+ City .location .nearest_neighbor (location , num + 1 ) & City .name .not_equals (city )
73
75
)
74
76
neighbors = qb .build ().find_with_scores ()
75
77
list_cities_with_scores (neighbors )
76
78
else :
77
79
print (f"no city found named '{ city } '" )
78
- except ValueError :
80
+ except ValueError :
79
81
print ("usage: city_neighbors <name>[,<num: default 5>]" )
80
-
82
+
81
83
def do_neighbors (self , args ):
82
84
"""find <num> neighbors next to geo-coord <lat> <long>.\n usage: neighbors <num>,<latitude>,<longitude>"""
83
85
try :
84
86
args = args .split (',' )
85
87
if len (args ) != 3 :
86
88
raise ValueError ()
87
89
num = int (args [0 ])
88
- geocoord = [ float (args [1 ]), float (args [2 ]) ]
90
+ geocoord = [float (args [1 ]), float (args [2 ])]
89
91
qb = self ._box .query (
90
92
City .location .nearest_neighbor (geocoord , num )
91
93
)
92
94
neighbors = qb .build ().find_with_scores ()
93
95
list_cities_with_scores (neighbors )
94
- except ValueError :
96
+ except ValueError :
95
97
print ("usage: neighbors <num>,<latitude>,<longitude>" )
96
-
98
+
97
99
def do_add (self , args : str ):
98
100
"""add new location\n usage: add <name>,<lat>,<long>"""
99
101
try :
@@ -105,11 +107,11 @@ def do_add(self, args: str):
105
107
long = float (args [2 ])
106
108
city = City ()
107
109
city .name = name
108
- city .location = [lat ,long ]
110
+ city .location = [lat , long ]
109
111
self ._box .put (city )
110
- except ValueError :
112
+ except ValueError :
111
113
print ("usage: add <name>,<latitude>,<longitude>" )
112
-
114
+
113
115
def do_exit (self , _ ):
114
116
"""close the program"""
115
117
raise SystemExit ()
0 commit comments