Skip to content

Commit edce194

Browse files
committed
Factor out _some to back one and one_or_zero
This changes the name of the NotFound exception to OutOfBounds, which is actually backwards incompatible with our documented interface. Fudgin' it. :/
1 parent 2d6b07d commit edce194

File tree

2 files changed

+56
-28
lines changed

2 files changed

+56
-28
lines changed

postgres.py

Lines changed: 41 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -194,14 +194,25 @@ def url_to_dsn(url):
194194
# Exceptions
195195
# ==========
196196

197-
class NotOne(Exception):
198-
def __init__(self, rowcount):
199-
self.rowcount = rowcount
197+
class OutOfBounds(Exception):
198+
199+
def __init__(self, n, lo, hi):
200+
self.n = n
201+
self.lo = lo
202+
self.hi = hi
203+
200204
def __str__(self):
201-
return "Got {0} rows instead of 1.".format(self.rowcount)
205+
msg = "Got {n} rows; expecting "
206+
if self.lo == self.hi:
207+
msg += "exactly {lo}."
208+
elif self.hi - self.lo == 1:
209+
msg += "{lo} or {hi}."
210+
else:
211+
msg += "between {lo} and {hi} (inclusive)."
212+
return msg.format(**self.__dict__)
202213

203-
class TooFew(NotOne): pass
204-
class TooMany(NotOne): pass
214+
class TooFew(OutOfBounds): pass
215+
class TooMany(OutOfBounds): pass
205216

206217

207218
# The Main Event
@@ -335,7 +346,7 @@ def one(self, sql, parameters=None, strict=None):
335346
By default, :py:attr:`strict` ends up evaluating to :py:class:`True`,
336347
in which case we raise :py:exc:`postgres.TooFew` or
337348
:py:exc:`postgres.TooMany` if the number of rows returned isn't exactly
338-
one (both are subclasses of :py:exc:`postgres.NotOne`). You can
349+
one (both are subclasses of :py:exc:`postgres.OutOfBounds`). You can
339350
override this behavior per-call with the :py:attr:`strict` argument
340351
here, or globally by passing :py:attr:`strict_one` to the
341352
:py:class:`~postgres.Postgres` constructor. If you use both, the
@@ -357,16 +368,14 @@ def one(self, sql, parameters=None, strict=None):
357368
else:
358369
strict = self.strict_one # user default
359370

360-
with self.get_cursor() as cursor:
361-
cursor.execute(sql, parameters)
362-
363-
if strict:
364-
if cursor.rowcount < 1:
365-
raise TooFew(cursor.rowcount)
366-
elif cursor.rowcount > 1:
367-
raise TooMany(cursor.rowcount)
371+
if strict:
372+
out = self._some(sql, parameters, 1, 1)
373+
else:
374+
with self.get_cursor() as cursor:
375+
cursor.execute(sql, parameters)
376+
out = cursor.fetchone()
377+
return out
368378

369-
return cursor.fetchone()
370379

371380
def one_or_zero(self, sql, parameters=None):
372381
"""Execute a query and return a single result or :py:class:`None`.
@@ -387,15 +396,24 @@ def one_or_zero(self, sql, parameters=None):
387396
No blam yet.
388397
389398
"""
390-
with self.get_cursor() as cursor:
391-
cursor.execute(sql, parameters)
399+
return self._some(sql, parameters, 0, 1)
400+
401+
402+
def _some(self, sql, parameters=None, lo=0, hi=1):
403+
404+
# This is undocumented (and largely untested) because I think it's a
405+
# rare case where this is wanted directly. It's here to make one and
406+
# one_or_zero DRY. Help yourself to it now that you've found it. :^)
407+
408+
with self.get_cursor() as txn:
409+
txn.execute(sql, parameters)
392410

393-
if cursor.rowcount < 0:
394-
raise TooFew(cursor.rowcount)
395-
elif cursor.rowcount > 1:
396-
raise TooMany(cursor.rowcount)
411+
if txn.rowcount < lo:
412+
raise TooFew(txn.rowcount, lo, hi)
413+
elif txn.rowcount > hi:
414+
raise TooMany(txn.rowcount, lo, hi)
397415

398-
return cursor.fetchone()
416+
return txn.fetchone()
399417

400418

401419
def get_cursor(self, *a, **kw):

tests.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
from postgres import Postgres, TooFew, TooMany
88
from psycopg2.extras import NamedTupleCursor
9-
from psycopg2 import InterfaceError
9+
from psycopg2 import InterfaceError, ProgrammingError
1010

1111

1212
DATABASE_URL = os.environ['DATABASE_URL']
@@ -89,14 +89,24 @@ def test_TooFew_message_is_helpful(self):
8989
self.db.one("SELECT * FROM foo WHERE bar='blah'", strict=True)
9090
except TooFew as exc:
9191
actual = str(exc)
92-
assert actual == "Got 0 rows instead of 1."
92+
assert actual == "Got 0 rows; expecting exactly 1."
9393

94-
def test_TooMany_message_is_helpful(self):
94+
def test_TooMany_message_is_helpful_for_two_options(self):
9595
try:
96-
self.db.one("SELECT * FROM foo", strict=True)
96+
self.db.one_or_zero("SELECT * FROM foo")
9797
except TooMany as exc:
9898
actual = str(exc)
99-
assert actual == "Got 2 rows instead of 1."
99+
assert actual == "Got 2 rows; expecting 0 or 1."
100+
101+
def test_TooMany_message_is_helpful_for_a_range(self):
102+
self.db.run("INSERT INTO foo VALUES ('blam')")
103+
self.db.run("INSERT INTO foo VALUES ('blim')")
104+
try:
105+
self.db._some("SELECT * FROM foo", lo=1, hi=3)
106+
except TooMany as exc:
107+
actual = str(exc)
108+
assert actual == \
109+
"Got 4 rows; expecting between 1 and 3 (inclusive)."
100110

101111

102112
class TestOne(WithData):

0 commit comments

Comments
 (0)