Convert fields names into valid Python identifiers in NamedTupleCursor

Close #211.
This commit is contained in:
Daniele Varrazzo 2018-01-29 02:41:44 +00:00
parent a84c9723b4
commit e5da79fcc8
3 changed files with 23 additions and 1 deletions

2
NEWS
View File

@ -7,6 +7,8 @@ What's new in psycopg 2.7.4
- Moving away from installing the wheel package by default.
Packages installed from wheel raise a warning on import. Added package
``psycopg2-binary`` to install from wheel instead (:ticket:`#543`).
- Convert fields names into valid Python identifiers in
`~psycopg2.extras.NamedTupleCursor` (:ticket:`#211`).
- Fixed Solaris 10 support (:ticket:`#532`).
- `cursor.mogrify()` can be called on closed cursors (:ticket:`#579`).
- Fixed setting session characteristics in corner cases on autocommit

View File

@ -368,7 +368,18 @@ class NamedTupleCursor(_cursor):
raise self._exc
else:
def _make_nt(self, namedtuple=namedtuple):
return namedtuple("Record", [d[0] for d in self.description or ()])
def f(s):
# NOTE: Python 3 actually allows unicode chars in fields
s = _re.sub('[^a-zA-Z0-9_]', '_', s)
# Python identifier cannot start with numbers, namedtuple fields
# cannot start with underscore. So...
if _re.match('^[0-9_]', s):
s = 'f' + s
return s
return namedtuple(
"Record", [f(d[0]) for d in self.description or ()])
class LoggingConnection(_connection):

View File

@ -382,6 +382,15 @@ class NamedTupleCursorTest(ConnectingTestCase):
curs.execute("update nttest set s = s")
self.assertRaises(psycopg2.ProgrammingError, curs.fetchall)
@skip_if_no_namedtuple
def test_bad_col_names(self):
curs = self.conn.cursor()
curs.execute('select 1 as "foo.bar_baz", 2 as "?column?", 3 as "3"')
rv = curs.fetchone()
self.assertEqual(rv.foo_bar_baz, 1)
self.assertEqual(rv.f_column_, 2)
self.assertEqual(rv.f3, 3)
@skip_if_no_namedtuple
def test_minimal_generation(self):
# Instrument the class to verify it gets called the minimum number of times.