Added errors.lookup() function

3 files changed, 18 insertions, 0 deletions

diff --git a/doc/src/errors.rst b/doc/src/errors.rst
index a4bd48a..061999b 100644
--- a/doc/src/errors.rst
+++ b/doc/src/errors.rst
@@ -56,3 +56,5 @@ idiomatic error handler:
 For completeness, the module also exposes all the DB-API-defined classes and
 :ref:`a few psycopg-specific exceptions <extension-exceptions>` previously
 exposed by the `!extensions` module.  One stop shop for all your mistakes...
+
+.. autofunction:: lookup
diff --git a/lib/errors.py b/lib/errors.py
index dbf829b..492d5aa 100644
--- a/lib/errors.py
+++ b/lib/errors.py
@@ -10,6 +10,14 @@ from psycopg2._psycopg import (
     QueryCanceledError, TransactionRollbackError)
 
 
+def lookup(code):
+    """Lookup an error code and return its exception class.
+
+    Raise `!KeyError` if the code is not found.
+    """
+    return _by_sqlstate[code]
+
+
 _by_sqlstate = {}
 
 
diff --git a/tests/test_errors.py b/tests/test_errors.py
index 5645e18..cb680a2 100755
--- a/tests/test_errors.py
+++ b/tests/test_errors.py
@@ -54,6 +54,14 @@ class ErrorsTests(ConnectingTestCase):
 
         self.assertEqual(type(e), self.conn.ProgrammingError)
 
+    def test_lookup(self):
+        from psycopg2 import errors
+
+        self.assertIs(errors.lookup('42P01'), errors.UndefinedTable)
+
+        with self.assertRaises(KeyError):
+            errors.lookup('XXXXX')
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)