From: Denis Laxalde <denis.laxalde@dalibo.com>
Date: Tue, 23 Feb 2021 17:52:12 +0000 (+0100)
Subject: Add namedtuple_row row factory
X-Git-Tag: 3.0.dev0~106^2~11
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=1a166783c2f339e6f5a064653e31ba4da16b1c85;p=thirdparty%2Fpsycopg.git

Add namedtuple_row row factory

Implementation for namedtuple class cache and fields name sanitizing is
taken from psycogp2.
---

diff --git a/docs/row-factories.rst b/docs/row-factories.rst
index eae335be3..1c620e7c0 100644
--- a/docs/row-factories.rst
+++ b/docs/row-factories.rst
@@ -58,3 +58,4 @@ Module `psycopg3.rows` contains available row factories:
 .. currentmodule:: psycopg3.rows
 
 .. autofunction:: dict_row
+.. autofunction:: namedtuple_row
diff --git a/psycopg3/psycopg3/rows.py b/psycopg3/psycopg3/rows.py
index c330b2673..18f7b93df 100644
--- a/psycopg3/psycopg3/rows.py
+++ b/psycopg3/psycopg3/rows.py
@@ -4,7 +4,10 @@ psycopg3 row factories
 
 # Copyright (C) 2021 The Psycopg Team
 
-from typing import Any, Callable, Dict, Sequence
+import functools
+import re
+from collections import namedtuple
+from typing import Any, Callable, Dict, Sequence, Tuple, Type
 
 from .cursor import BaseCursor
 from .proto import ConnectionType
@@ -21,3 +24,37 @@ def dict_row(
         return dict(zip(titles, values))
 
     return make_row
+
+
+def namedtuple_row(
+    cursor: BaseCursor[ConnectionType],
+) -> Callable[[Sequence[Any]], Tuple[Any, ...]]:
+    """Row factory to represent rows as `~collections.namedtuple`."""
+
+    def make_row(values: Sequence[Any]) -> Tuple[Any, ...]:
+        assert cursor.description
+        key = tuple(c.name for c in cursor.description)
+        nt = _make_nt(key)
+        rv = nt._make(values)  # type: ignore[attr-defined]
+        return rv  # type: ignore[no-any-return]
+
+    return make_row
+
+
+# ascii except alnum and underscore
+_re_clean = re.compile(
+    "[" + re.escape(" !\"#$%&'()*+,-./:;<=>?@[\\]^`{|}~") + "]"
+)
+
+
+@functools.lru_cache(512)
+def _make_nt(key: Sequence[str]) -> Type[Tuple[Any, ...]]:
+    fields = []
+    for s in key:
+        s = _re_clean.sub("_", s)
+        # Python identifier cannot start with numbers, namedtuple fields
+        # cannot start with underscore. So...
+        if s[0] == "_" or "0" <= s[0] <= "9":
+            s = "f" + s
+        fields.append(s)
+    return namedtuple("Row", fields)
diff --git a/tests/test_rows.py b/tests/test_rows.py
index d130965be..06168e698 100644
--- a/tests/test_rows.py
+++ b/tests/test_rows.py
@@ -11,3 +11,34 @@ def test_dict_row(conn):
     assert cur.nextset()
     assert cur.fetchall() == [{"number": 1}]
     assert not cur.nextset()
+
+
+def test_namedtuple_row(conn):
+    cur = conn.cursor(row_factory=rows.namedtuple_row)
+    cur.execute("select 'bob' as name, 3 as id")
+    (person1,) = cur.fetchall()
+    assert f"{person1.name} {person1.id}" == "bob 3"
+
+    ci1 = rows._make_nt.cache_info()
+    assert ci1.hits == 0 and ci1.misses == 1
+
+    cur.execute("select 'alice' as name, 1 as id")
+    (person2,) = cur.fetchall()
+    assert type(person2) is type(person1)
+
+    ci2 = rows._make_nt.cache_info()
+    assert ci2.hits == 1 and ci2.misses == 1
+
+    cur.execute("select 'foo', 1 as id")
+    (r0,) = cur.fetchall()
+    assert r0.f_column_ == "foo"
+    assert r0.id == 1
+
+    cur.execute("select 'a' as letter; select 1 as number")
+    (r1,) = cur.fetchall()
+    assert r1.letter == "a"
+    assert cur.nextset()
+    (r2,) = cur.fetchall()
+    assert r2.number == 1
+    assert not cur.nextset()
+    assert type(r1) is not type(r2)