Skip to content

Commit

Permalink
Better PyPy support.
Browse files Browse the repository at this point in the history
  • Loading branch information
jamadden committed Oct 24, 2023
1 parent ca851c3 commit c1c8248
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 17 deletions.
1 change: 1 addition & 0 deletions .pylintrc
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,7 @@ disable=wrong-import-position,
differing-type-doc,
compare-to-zero,
docstring-first-line-empty,
too-many-try-statements,

enable=consider-using-augmented-assign

Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
'zope.site',
'zope.testrunner',
'testgres',
'psycopg2-binary; python_implementation != "PyPy"',
]

def _read(fname):
Expand Down
50 changes: 33 additions & 17 deletions src/nti/testing/layers/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
The APIs are preliminary and may change.
This is only supported on platforms that can install ``psycopg2``.
"""
from contextlib import contextmanager
import functools
Expand All @@ -18,9 +20,23 @@
import unittest
from unittest.mock import patch

import psycopg2
import psycopg2.extras
import psycopg2.pool
#import psycopg2
#import psycopg2.extras
#import psycopg2.pool

try:
from psycopg2 import ProgrammingError
except ImportError:
ThreadedConnectionPool = None
DictCursor = None
class IntegrityError(Exception):
"""Never thrown"""
ProgrammingError = InternalError = IntegrityError
else:
from psycopg2.pool import ThreadedConnectionPool
from psycopg2.extras import DictCursor
from psycopg2 import IntegrityError
from psycopg2 import InternalError

import testgres

Expand Down Expand Up @@ -50,11 +66,11 @@
# NTI_SAVE_DB is either 1/on/true (case-insensitive)
# or a file name.
val = os.environ['NTI_SAVE_DB']
if val.lower() in ('0', 'off', 'false', 'no'):
if val.lower() in {'0', 'off', 'false', 'no'}:
SAVE_DATABASE_ON_TEARDOWN = False
else:
SAVE_DATABASE_ON_TEARDOWN = True
if val.lower() not in ('1', 'on', 'true', 'yes'):
if val.lower() not in {'1', 'on', 'true', 'yes'}:
SAVE_DATABASE_FILENAME = val


Expand Down Expand Up @@ -117,7 +133,7 @@ class DatabaseLayer(object):
connection_pool = None


connection_pool_klass = psycopg2.pool.ThreadedConnectionPool
connection_pool_klass = ThreadedConnectionPool
connection_pool_minconn = 1
connection_pool_maxconn = 51

Expand Down Expand Up @@ -195,7 +211,7 @@ def setUp(cls):
dbname=cls.DATABASE_NAME,
host='localhost',
port=cls.postgres_node.port,
cursor_factory=psycopg2.extras.DictCursor
cursor_factory=DictCursor,
)

cls.postgres_dsn = "host=%s dbname=%s port=%s" % (
Expand Down Expand Up @@ -271,14 +287,14 @@ def borrowed_connection(cls):

@classmethod
def truncate_table(cls, conn, table_name):
"Transactionally truncate the given *table_name* using *conn*"
"""Transactionally truncate the given *table_name* using *conn*"""
try:

with conn.cursor() as cur:
cur.execute(
'TRUNCATE TABLE ' + table_name + ' CASCADE'
)
except (psycopg2.ProgrammingError, psycopg2.InternalError):
except (ProgrammingError, InternalError):
# Table doesn't exist, not a full schema,
# ignore.
# OR:
Expand All @@ -293,7 +309,7 @@ def truncate_table(cls, conn, table_name):

@classmethod
def drop_relation(cls, relation, kind='TABLE', idempotent=False):
"Drops the *relation* of type *kind* (default table), in new transaction."
"""Drops the *relation* of type *kind* (default table), in new transaction."""
with cls.borrowed_connection() as conn:
with conn.cursor() as cur:
if idempotent:
Expand Down Expand Up @@ -371,9 +387,9 @@ def print_size_report(cls):
print()
fmt = "| {table_name:35s} | {total:10s} | {index:10s} | {toast:10s} | {table:10s}"
for row in rows:
if not extra_query and row['total'] in (
'72 kB', '32 kB', '24 kB', '16 kB', '8192 bytes'
):
if not extra_query and row['total'] in {
'72 kB', '32 kB', '24 kB', '16 kB', '8192 bytes'
}:
continue
print(fmt.format(
**{k: v if v else '<null>' for k, v in row.items()}
Expand Down Expand Up @@ -453,7 +469,7 @@ def _tangle_schema_if_needed(cls):
ex = e
except subprocess.CalledProcessError as e:
output = ex.output
ex = e
ex = e # pylint:disable=redefined-variable-type

output = output.decode('utf-8')

Expand Down Expand Up @@ -533,7 +549,7 @@ def push(cls, layer):
dbname=layer.DATABASE_NAME,
host='localhost',
port=new_node.port,
cursor_factory=psycopg2.extras.DictCursor
cursor_factory=DictCursor
)

@classmethod
Expand Down Expand Up @@ -644,10 +660,10 @@ class DatabaseTestCase(unittest.TestCase):
@contextmanager
def assertRaisesIntegrityError(self, match=None):
if match:
with self.assertRaisesRegex(psycopg2.IntegrityError, match) as exc:
with self.assertRaisesRegex(IntegrityError, match) as exc:
yield exc
else:
with self.assertRaises(psycopg2.IntegrityError) as exc:
with self.assertRaises(IntegrityError) as exc:
yield exc

# We can't do any queries after an error is raised
Expand Down

0 comments on commit c1c8248

Please sign in to comment.