From 5844e989c43b6c25363e21cf3ffe310d3d94d52b Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 26 Sep 2010 16:57:07 +0100
Subject: [PATCH] Added function to parse an hstore into a dict.

---
 lib/extras.py         | 52 ++++++++++++++++++++++++++++++++++++++++++-
 tests/types_extras.py | 27 ++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/lib/extras.py b/lib/extras.py
index a402d11a..4b20bb64 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -35,7 +35,7 @@ try:
 except:
     logging = None
 
-from psycopg2 import DATETIME, DataError
+from psycopg2 import DATETIME, DataError, InterfaceError
 from psycopg2 import extensions as _ext
 from psycopg2.extensions import cursor as _cursor
 from psycopg2.extensions import connection as _connection
@@ -515,5 +515,55 @@ class HstoreAdapter(object):
         v.prepare(self.conn)
         return "hstore(%s, %s)" % (k.getquoted(), v.getquoted())
 
+_re_hstore = regex.compile(r"""
+    # hstore key:
+    "( # catch in quotes
+        (?: # many of
+            [^"] # either not a quote
+            # or a quote escaped, i.e. preceded by an odd number of backslashes
+            | [^\\] (?:\\\\)* \\"
+        )*
+    )"
+    \s*=>\s* # hstore value
+    (?:
+        NULL # the value can be null - not catched
+        # or the same quoted string of the key
+        | "((?:[^"] | [^\\] (?:\\\\)* \\" )*)"
+    )
+    (?:\s*,\s*|$) # pairs separated by comma or end of string.
+""", regex.VERBOSE)
+
+def parse_hstore(s, cur):
+    """Parse an hstore representation in a Python string.
+
+    The hstore is represented as something like::
+
+        "a"=>"1", "b"=>"2"
+
+    with backslash-escaped strings.
+    """
+    if s is None:
+        return None
+
+    rv = {}
+    start = 0
+    for m in _re_hstore.finditer(s):
+        if m is None or m.start() != start:
+            raise InterfaceError(
+                "error parsing hstore pair at char %d" % start)
+        k = m.group(1).decode("string_escape")
+        v = m.group(2)
+        if v is not None:
+            v = v.decode("string_escape")
+
+        rv[k] = v
+        start = m.end()
+
+    if start < len(s):
+        raise InterfaceError(
+            "error parsing hstore: unparsed data after char %d" % start)
+
+    return rv
+
 
 __all__ = filter(lambda k: not k.startswith('_'), locals().keys())
diff --git a/tests/types_extras.py b/tests/types_extras.py
index c59479a7..559e6950 100644
--- a/tests/types_extras.py
+++ b/tests/types_extras.py
@@ -157,6 +157,33 @@ class HstoreTestCase(unittest.TestCase):
         encc = u'\xe0'.encode(psycopg2.extensions.encodings[self.conn.encoding])
         self.assertEqual(ii[3], ("E'd'", "E'%s'" % encc))
 
+    def test_parse(self):
+        from psycopg2.extras import parse_hstore
+
+        def ok(s, d):
+            self.assertEqual(parse_hstore(s, None), d)
+
+        ok(None, None)
+        ok('', {})
+        ok('"a"=>"1", "b"=>"2"', {'a': '1', 'b': '2'})
+        ok('"a"  => "1" ,"b"  =>  "2"', {'a': '1', 'b': '2'})
+        ok('"a"=>NULL, "b"=>"2"', {'a': None, 'b': '2'})
+        ok('"a"=>"\'", "\'"=>"2"', {'a': "'", "'": '2'})
+        ok('"a"=>"1", "b"=>NULL', {'a': '1', 'b': None})
+        ok(r'"a\\"=>"1"', {'a\\': '1'})
+        ok(r'"a\""=>"1"', {'a"': '1'})
+        ok(r'"a\\\""=>"1"', {r'a\"': '1'})
+        ok(r'"a\\\\\""=>"1"', {r'a\\"': '1'})
+
+        def ko(s):
+            self.assertRaises(psycopg2.InterfaceError, parse_hstore, s, None)
+
+        ko('a')
+        ko('"a"')
+        ko(r'"a\\""=>"1"')
+        ko(r'"a\\\\""=>"1"')
+        ko('"a=>"1"')
+        ko('"a"=>"1", "b"=>NUL')
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)