add possibility to don't replace non-words by default (#361)

Add possibility to don't replace non-words by default

clickhouse_driver/client.py

@@ -438,7 +438,7 @@ def execute_iter(
 
     def query_dataframe(
             self, query, params=None, external_tables=None, query_id=None,
-            settings=None):
+            settings=None, replace_nonwords=True):
         """
         *New in version 0.2.0.*
 
@@ -453,6 +453,7 @@ def query_dataframe(
                          ClickHouse server will generate it.
         :param settings: dictionary of query settings.
                          Defaults to ``None`` (no additional settings).
+        :param replace_nonwords: boolean to replace non-words in column names to underscores
         :return: pandas DataFrame.
         """
 
@@ -467,7 +468,9 @@ def query_dataframe(
             settings=settings
         )
 
-        columns = [re.sub(r'\W', '_', name) for name, type_ in columns]
+        if replace_nonwords:
+            columns = [re.sub(r'\W', '_', name) for name, type_ in columns]
+
         return pd.DataFrame(
             {col: d for d, col in zip(data, columns)}, columns=columns
         )

tests/columns/test_columns_names.py

@@ -0,0 +1,17 @@
+from tests.testcase import BaseTestCase
+
+class ColumnsNamesTestCase(BaseTestCase):
+
+    def test_columns_names_replace_nonwords(self):
+        columns = (
+            'regular Int64, '
+            'CamelCase Int64, '
+            'With_Underscores Int64, '
+            'Any%different.Column? Int64'
+        )
+
+        expected_columns = ['regular', 'CamelCase', 'With_Underscores', 'Any%different.Column?']
+
+        with self.create_table(columns):
+            df = self.client.query_dataframe('SELECT * FROM test', replace_nonwords=False)
+            self.assertTrue(expected_columns.equals(list(df.columns)))