From c1c824852b3803e54c653a36c0c3e49d1c413e24 Mon Sep 17 00:00:00 2001
From: Jason Madden <jamadden@gmail.com>
Date: Tue, 24 Oct 2023 12:03:37 -0500
Subject: [PATCH] Better PyPy support.

---
 .pylintrc                          |  1 +
 setup.py                           |  1 +
 src/nti/testing/layers/postgres.py | 50 ++++++++++++++++++++----------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/.pylintrc b/.pylintrc
index a2d4197..5dee94c 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -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
 
diff --git a/setup.py b/setup.py
index 7d6cf58..feb6528 100755
--- a/setup.py
+++ b/setup.py
@@ -14,6 +14,7 @@
     'zope.site',
     'zope.testrunner',
     'testgres',
+    'psycopg2-binary; python_implementation != "PyPy"',
 ]
 
 def _read(fname):
diff --git a/src/nti/testing/layers/postgres.py b/src/nti/testing/layers/postgres.py
index 717c83d..be1ad79 100644
--- a/src/nti/testing/layers/postgres.py
+++ b/src/nti/testing/layers/postgres.py
@@ -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
@@ -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
 
@@ -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
 
 
@@ -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
 
@@ -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" %  (
@@ -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:
@@ -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:
@@ -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()}
@@ -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')
 
@@ -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
@@ -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