From 29007fd516814dc8b10c518a282b625bc024e3a9 Mon Sep 17 00:00:00 2001
From: Tyler Hobbs <tyler@datastax.com>
Date: Fri, 9 May 2014 15:12:31 -0500
Subject: [PATCH] Don't strip trailing _ in column names

Fixes PYTHON-56:
---
 CHANGELOG.rst            |  9 +++++++++
 cassandra/decoder.py     | 14 +++++++++++---
 tests/unit/test_types.py |  5 +++--
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2e567cef..234b16ba 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,3 +1,12 @@
+1.1.3
+=====
+In Progress
+
+Bug Fixes
+---------
+* Don't strip trailing underscores from column names when using the
+  named_tuple_factory (PYTHON-56)
+
 1.1.2
 =====
 May 8, 2014
diff --git a/cassandra/decoder.py b/cassandra/decoder.py
index 9893c12d..0a014bd3 100644
--- a/cassandra/decoder.py
+++ b/cassandra/decoder.py
@@ -63,12 +63,20 @@ HEADER_DIRECTION_TO_CLIENT = 0x80
 HEADER_DIRECTION_MASK = 0x80
 
 
-NON_ALPHA_REGEX = re.compile('\W')
-END_UNDERSCORE_REGEX = re.compile('^_*(\w*[a-zA-Z0-9])_*$')
+NON_ALPHA_REGEX = re.compile('[^a-zA-Z0-9]')
+START_BADCHAR_REGEX = re.compile('^[^a-zA-Z0-9]*')
+END_BADCHAR_REGEX = re.compile('[^a-zA-Z0-9_]*$')
+
+_clean_name_cache = {}
 
 
 def _clean_column_name(name):
-    return END_UNDERSCORE_REGEX.sub("\g<1>", NON_ALPHA_REGEX.sub("_", name))
+    try:
+        return _clean_name_cache[name]
+    except KeyError:
+        clean = NON_ALPHA_REGEX.sub("_", START_BADCHAR_REGEX.sub("", END_BADCHAR_REGEX.sub("", name)))
+        _clean_name_cache[name] = clean
+        return clean
 
 
 def tuple_factory(colnames, rows):
diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py
index 323553b5..854f637b 100644
--- a/tests/unit/test_types.py
+++ b/tests/unit/test_types.py
@@ -130,13 +130,14 @@ class TypeTests(unittest.TestCase):
         self.assertEqual(cql_typename('org.apache.cassandra.db.marshal.ListType(IntegerType)'), 'list<varint>')
 
     def test_named_tuple_colname_substitution(self):
-        colnames = ("func(abc)", "[applied]", "func(func(abc))", "foo_bar")
-        rows = [(1, 2, 3, 4)]
+        colnames = ("func(abc)", "[applied]", "func(func(abc))", "foo_bar", "foo_bar_")
+        rows = [(1, 2, 3, 4, 5)]
         result = named_tuple_factory(colnames, rows)[0]
         self.assertEqual(result[0], result.func_abc)
         self.assertEqual(result[1], result.applied)
         self.assertEqual(result[2], result.func_func_abc)
         self.assertEqual(result[3], result.foo_bar)
+        self.assertEqual(result[4], result.foo_bar_)
 
     def test_parse_casstype_args(self):
         class FooType(CassandraType):