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 eae335be323b35705177cacab4093722be2aed5b..1c620e7c0900def9fb15b378d26c8cba0ea56c81 100644 (file)
--- 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 c330b26738d6021ee8b9363a4aba4a5be4408890..18f7b93df28a174a7fb0a0e677cad26fba50a10d 100644 (file)
--- 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 d130965be22f5843f0c5003b00f37a45c58926ec..06168e69834a66fea5c9938b62c373ddffa3e3ef 100644 (file)
--- 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)