Skip to content

Commit cf8afb8

Browse files
committed
vectorsearch-cities example: use explicit imports and fix formatting
1 parent ca79c96 commit cf8afb8

File tree

1 file changed

+24
-22
lines changed

1 file changed

+24
-22
lines changed

example/vectorsearch-cities/main.py

+24-22
Original file line numberDiff line numberDiff line change
@@ -1,52 +1,54 @@
11
from cmd import Cmd
2-
from objectbox import *
2+
from objectbox import Entity, Float32Vector, HnswIndex, Id, Store, String
33
import time
44
import csv
55
import os
66

7+
78
@Entity()
89
class City:
910
id = Id()
1011
name = String()
11-
location = Float32Vector(index=HnswIndex(
12-
dimensions=2,
13-
distance_type=VectorDistanceType.EUCLIDEAN
14-
))
12+
location = Float32Vector(index=HnswIndex(dimensions=2))
13+
1514

1615
def list_cities(cities):
1716
print("{:3s} {:25s} {:>9s} {:>9s}".format("ID", "Name", "Latitude", "Longitude"))
1817
for city in cities:
1918
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+
2121

2222
def list_cities_with_scores(city_score_tuples):
2323
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:
2525
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+
2728

2829
class VectorSearchCitiesCmd(Cmd):
2930
prompt = "> "
31+
3032
def __init__(self, *args):
3133
Cmd.__init__(self, *args)
3234
dbdir = "cities-db"
3335
new_db = not os.path.exists(dbdir)
3436
self._store = Store(directory=dbdir)
3537
self._box = self._store.box(City)
36-
if new_db:
38+
if new_db:
3739
with open(os.path.join(os.path.dirname(__file__), 'cities.csv')) as f:
38-
r = csv.reader(f)
40+
r = csv.reader(f)
3941
cities = []
4042
for row in r:
4143
city = City()
4244
city.name = row[0]
43-
city.location = [ row[1], row[2] ]
45+
city.location = [row[1], row[2]]
4446
cities.append(city)
4547
self._box.put(*cities)
4648

4749
def do_ls(self, name: str = ""):
4850
"""list all cities or starting with <prefix>\nusage: ls [<prefix>]"""
49-
qb = self._box.query( City.name.starts_with(name) )
51+
qb = self._box.query(City.name.starts_with(name))
5052
query = qb.build()
5153
list_cities(query.find())
5254

@@ -62,38 +64,38 @@ def do_city_neighbors(self, args: str):
6264
num = 5
6365
if len(args) == 2:
6466
num = int(args[1])
65-
qb = self._box.query( City.name.equals(city) )
67+
qb = self._box.query(City.name.equals(city))
6668
query = qb.build()
6769
cities = query.find()
6870
if len(cities) == 1:
6971
location = cities[0].location
7072
# +1 for the city
7173
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)
7375
)
7476
neighbors = qb.build().find_with_scores()
7577
list_cities_with_scores(neighbors)
7678
else:
7779
print(f"no city found named '{city}'")
78-
except ValueError:
80+
except ValueError:
7981
print("usage: city_neighbors <name>[,<num: default 5>]")
80-
82+
8183
def do_neighbors(self, args):
8284
"""find <num> neighbors next to geo-coord <lat> <long>.\nusage: neighbors <num>,<latitude>,<longitude>"""
8385
try:
8486
args = args.split(',')
8587
if len(args) != 3:
8688
raise ValueError()
8789
num = int(args[0])
88-
geocoord = [ float(args[1]), float(args[2]) ]
90+
geocoord = [float(args[1]), float(args[2])]
8991
qb = self._box.query(
9092
City.location.nearest_neighbor(geocoord, num)
9193
)
9294
neighbors = qb.build().find_with_scores()
9395
list_cities_with_scores(neighbors)
94-
except ValueError:
96+
except ValueError:
9597
print("usage: neighbors <num>,<latitude>,<longitude>")
96-
98+
9799
def do_add(self, args: str):
98100
"""add new location\nusage: add <name>,<lat>,<long>"""
99101
try:
@@ -105,11 +107,11 @@ def do_add(self, args: str):
105107
long = float(args[2])
106108
city = City()
107109
city.name = name
108-
city.location = [lat,long]
110+
city.location = [lat, long]
109111
self._box.put(city)
110-
except ValueError:
112+
except ValueError:
111113
print("usage: add <name>,<latitude>,<longitude>")
112-
114+
113115
def do_exit(self, _):
114116
"""close the program"""
115117
raise SystemExit()

0 commit comments

Comments
 (0)