Added caching of types generated by NamedTupleCursor

see #838

NEWS

@@ -30,6 +30,7 @@ New features:
   maintain columns order (:ticket:`#177`).
 - Added `~psycopg2.extensions.Diagnostics.severity_nonlocalized` attribute on
   the `~psycopg2.extensions.Diagnostics` object (:ticket:`#783`).
+- More efficient `~psycopg2.extras.NamedTupleCursor` (:ticket:`#838`).
 
 Other changes:
 

lib/extras.py

@@ -330,6 +330,7 @@ class NamedTupleCursor(_cursor):
         "abc'def"
     """
     Record = None
+    MAX_CACHE = 1024
 
     def execute(self, query, vars=None):
         self.Record = None
@@ -381,21 +382,31 @@ class NamedTupleCursor(_cursor):
         except StopIteration:
             return
 
-    def _make_nt(self):
-        # ascii except alnum and underscore
-        nochars = ' !"#$%&\'()*+,-./:;<=>?@[\\]^`{|}~'
-        re_clean = _re.compile('[' + _re.escape(nochars) + ']')
+    # ascii except alnum and underscore
+    _re_clean = _re.compile(
+        '[' + _re.escape(' !"#$%&\'()*+,-./:;<=>?@[\\]^`{|}~') + ']')
 
-        def f(s):
-            s = re_clean.sub('_', s)
-            # Python identifier cannot start with numbers, namedtuple fields
-            # cannot start with underscore. So...
+    _nt_cache = OrderedDict()
+
+    def _make_nt(self):
+        key = tuple(d[0] for d in (self.description or ()))
+        nt = self._nt_cache.get(key)
+        if nt is not None:
+            return nt
+
+        fields = []
+        for s in key:
+            s = self._re_clean.sub('_', s)
             if s[0] == '_' or '0' <= s[0] <= '9':
                 s = 'f' + s
+            fields.append(s)
 
-            return s
+        nt = namedtuple("Record", fields)
+        self._nt_cache[key] = nt
+        while len(self._nt_cache) > self.MAX_CACHE:
+            self._nt_cache.popitem(last=False)
 
-        return namedtuple("Record", [f(d[0]) for d in self.description or ()])
+        return nt
 
 
 class LoggingConnection(_connection):

tests/test_extras_dictcursor.py

@@ -578,6 +578,42 @@ class NamedTupleCursorTest(ConnectingTestCase):
         for i, t in enumerate(curs):
             self.assertEqual(i + 1, curs.rownumber)
 
+    def test_cache(self):
+        curs = self.conn.cursor()
+        curs.execute("select 10 as a, 20 as b")
+        r1 = curs.fetchone()
+        curs.execute("select 10 as a, 20 as c")
+        r2 = curs.fetchone()
+        curs.execute("select 10 as a, 30 as b")
+        r3 = curs.fetchone()
+
+        self.assert_(type(r1) is type(r3))
+        self.assert_(type(r1) is not type(r2))
+
+    def test_max_cache(self):
+        from psycopg2.extras import NamedTupleCursor
+        old_max_cache = NamedTupleCursor.MAX_CACHE
+        NamedTupleCursor.MAX_CACHE = 10
+        try:
+            NamedTupleCursor._nt_cache.clear()
+            curs = self.conn.cursor()
+            for i in range(10):
+                curs.execute("select 1 as f%s" % i)
+                curs.fetchone()
+
+            self.assertEqual(len(NamedTupleCursor._nt_cache), 10)
+            for i in range(10):
+                self.assert_(('f%s' % i,) in NamedTupleCursor._nt_cache)
+
+            curs.execute("select 1 as f10")
+            curs.fetchone()
+            self.assertEqual(len(NamedTupleCursor._nt_cache), 10)
+            self.assert_(('f10',) in NamedTupleCursor._nt_cache)
+            self.assert_(('f0',) not in NamedTupleCursor._nt_cache)
+
+        finally:
+            NamedTupleCursor.MAX_CACHE = old_max_cache
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)