Added implementation of python dict adapter to hstore.

lib/extras.py

@@ -473,4 +473,47 @@ def wait_select(conn):
             raise OperationalError("bad state from poll: %s" % state)
 
 
+class HstoreAdapter(object):
+    """Adapt a Python dict to the hstore syntax."""
+    def __init__(self, wrapped):
+        self.wrapped = wrapped
+
+    def prepare(self, conn):
+        self.conn = conn
+
+        # select the most efficient getquoted implementation
+        if conn.server_version >= 90000:
+            self.getquoted = self._getquoted_9
+
+    def _getquoted_8(self):
+        """Use the operators available in PG pre-9.0."""
+        adapt = _ext.adapt
+        rv = []
+        for k, v in self.wrapped.iteritems():
+            k = adapt(k)
+            k.prepare(self.conn)
+            k = k.getquoted()
+
+            if v is not None:
+                v = adapt(v)
+                v.prepare(self.conn)
+                v = v.getquoted()
+            else:
+                v = 'NULL'
+
+            rv.append("(%s => %s)" % (k, v))
+
+        return "(" + '||'.join(rv) + ")"
+
+    getquoted = _getquoted_8
+
+    def _getquoted_9(self):
+        """Use the hstore(text[], text[]) function."""
+        k = _ext.adapt(self.wrapped.keys())
+        k.prepare(self.conn)
+        v = _ext.adapt(self.wrapped.values())
+        v.prepare(self.conn)
+        return "hstore(%s, %s)" % (k.getquoted(), v.getquoted())
+
+
 __all__ = filter(lambda k: not k.startswith('_'), locals().keys())

tests/types_extras.py

@@ -18,8 +18,10 @@ try:
     import decimal
 except:
     pass
+import re
 import sys
 import unittest
+import warnings
 
 import psycopg2
 import psycopg2.extras
@@ -101,6 +103,61 @@ class TypesExtrasTests(unittest.TestCase):
         except psycopg2.ProgrammingError, err:
             self.failUnless(str(err) == "can't adapt type 'Foo'")
 
+
+class HstoreTestCase(unittest.TestCase):
+    def setUp(self):
+        self.conn = psycopg2.connect(tests.dsn)
+
+    def test_adapt_8(self):
+        if self.conn.server_version >= 90000:
+            warnings.warn("skipping dict adaptation with PG pre-9 syntax")
+            return
+
+        from psycopg2.extras import HstoreAdapter
+
+        o = {'a': '1', 'b': "'", 'c': None, 'd': u'\xe0'}
+        a = HstoreAdapter(o)
+        a.prepare(self.conn)
+        q = a.getquoted()
+
+        self.assert_(q.startswith("(("), q)
+        self.assert_(q.endswith("))"), q)
+        ii = q[1:-1].split("||")
+        ii.sort()
+
+        self.assertEqual(ii[0], "(E'a' => E'1')")
+        self.assertEqual(ii[1], "(E'b' => E'''')")
+        self.assertEqual(ii[2], "(E'c' => NULL)")
+        encc = u'\xe0'.encode(psycopg2.extensions.encodings[self.conn.encoding])
+        self.assertEqual(ii[3], "(E'd' => E'%s')" % encc)
+
+    def test_adapt_9(self):
+        if self.conn.server_version < 90000:
+            warnings.warn("skipping dict adaptation with PG 9 syntax")
+            return
+
+        from psycopg2.extras import HstoreAdapter
+
+        o = {'a': '1', 'b': "'", 'c': None, 'd': u'\xe0'}
+        a = HstoreAdapter(o)
+        a.prepare(self.conn)
+        q = a.getquoted()
+
+        m = re.match(r'hstore\(ARRAY\[([^\]]+)\], ARRAY\[([^\]]+)\]\)', q)
+        self.assert_(m, repr(q))
+
+        kk = m.group(1).split(", ")
+        vv = m.group(2).split(", ")
+        ii = zip(kk, vv)
+        ii.sort()
+
+        self.assertEqual(ii[0], ("E'a'", "E'1'"))
+        self.assertEqual(ii[1], ("E'b'", "E''''"))
+        self.assertEqual(ii[2], ("E'c'", "NULL"))
+        encc = u'\xe0'.encode(psycopg2.extensions.encodings[self.conn.encoding])
+        self.assertEqual(ii[3], ("E'd'", "E'%s'" % encc))
+
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)