Add possibility for few SQL variables in Python file

Python file reading tests became slower due to this. Now there is some extra checking of variables, that may be the reson.

However, it seems that SQL file reading tests became slower as well even tough their part of code was not changed.

Author: AzisK

readsql/__main__.py

@@ -33,15 +33,18 @@ def read_file(file_name, inplace=True):
             out.write(lines)
 
 
-def read_python_file(file_name, variable='query', inplace=True):
+def read_python_file(file_name, variables=None, inplace=True):
+    variables = variables if variables else ['query']
+    variables_regex = f"(?:{'|'.join(variables)})" if len(variables) > 1 else variables[0]
+
     with open(file_name, 'r') as inp:
         lines = inp.read()
         subs = []
 
         regex = [
             m
             for m in re.finditer(
-                r'(?:\s*' + variable + r'\s*=\s*f?)(?:"{1,3}|\'{1,3})([^"]*)(:?"|\')',
+                r'(?:\s*' + variables_regex + r'\s*=\s*f?)(?:"{1,3}|\'{1,3})([^"]*)(:?"|\')',
                 lines,
             )
         ]

readsql/parse_args.py

@@ -13,8 +13,8 @@ def parse_args():
     parser.add_argument(
         '-py',
         '--python_var',
-        action='store',
-        default='query',
+        nargs='+',
+        default=['query'],
         help='Look for certain variables in a Python file (default variable is query)',
     )
 

tests/speed.txt

@@ -1,15 +1,17 @@
-func:test_read_python_file__var:(((), {'variable': 'sql'}))__took: 0.000690 sec
+func:test_read_python_file__var:(((), {'variable': ['sql']}))__took: 0.001723 sec
 
-func:test_read_file__var:(((), {}))__took: 0.002316 sec
+func:test_read_file__var:(((), {}))__took: 0.005423 sec
 
-func:test_read_python_file__var:(((), {}))__took: 0.000885 sec
+func:test_read_python_file__var:(((), {}))__took: 0.003344 sec
 
-func:test_double_select__var:(((), {}))__took: 0.000245 sec
+func:test_double_select__var:(((), {}))__took: 0.000642 sec
 
-func:test_select_from_groub_by_where__var:(((), {}))__took: 0.000202 sec
+func:test_select_from_groub_by_where__var:(((), {}))__took: 0.000595 sec
 
-func:test_is_not_null__var:(((), {}))__took: 0.000236 sec
+func:test_is_not_null__var:(((), {}))__took: 0.000583 sec
 
-func:test_distinct__var:(((), {}))__took: 0.000196 sec
+func:test_distinct__var:(((), {}))__took: 0.001059 sec
 
-func:test_create_table_if_not_exists__var:(((), {}))__took: 0.000252 sec
+func:test_create_table_if_not_exists__var:(((), {}))__took: 0.000418 sec
+
+func:test_read_python_file__var:(((), {'variable': ['sql', 'query_template', 'query']}))__took: 0.001387 sec

tests/sql_in_python_variables_example.py

@@ -0,0 +1,16 @@
+def get_sql():
+    limit = 6
+    sql = f"SELEct speed from world where animal='dolphin' limit {limit}"
+    return sql
+
+
+def get_query_template():
+    limit = 6
+    query_template = f"SELEct speed from world where animal='dolphin' group by family limit {limit}"
+    return query_template
+
+
+def get_query():
+    limit = 99
+    query = f"SELEct speed from world where animal='dolphin' and name is not null group by family limit {limit}"
+    return query

tests/sql_in_python_variables_example_correct.py

@@ -0,0 +1,16 @@
+def get_sql():
+    limit = 6
+    sql = f"SELECT speed FROM world WHERE animal='dolphin' LIMIT {limit}"
+    return sql
+
+
+def get_query_template():
+    limit = 6
+    query_template = f"SELECT speed FROM world WHERE animal='dolphin' GROUP BY family LIMIT {limit}"
+    return query_template
+
+
+def get_query():
+    limit = 99
+    query = f"SELECT speed FROM world WHERE animal='dolphin' and name IS NOT NULL GROUP BY family LIMIT {limit}"
+    return query

tests/test_files.py

@@ -29,8 +29,18 @@ def test_read_python_file():
 def test_read_python_file_variable_wrap():
     @timing
     def test_read_python_file(variable):
-        return rsql.read_python_file(file_name=DIR + '/sql_in_python_variable_example.py', variable=variable, inplace=False)
+        return rsql.read_python_file(file_name=DIR + '/sql_in_python_variable_example.py', variables=variable, inplace=False)
 
-    example = test_read_python_file(variable='sql')
+    example = test_read_python_file(variable=['sql'])
     with open(DIR + '/sql_in_python_variable_example_correct.py', 'r') as inp:
         assert inp.read() == example
+
+
+def test_read_python_file_variables_wrap():
+    @timing
+    def test_read_python_file(variable):
+        return rsql.read_python_file(file_name=DIR + '/sql_in_python_variables_example.py', variables=variable, inplace=False)
+
+    example = test_read_python_file(variable=['sql', 'query_template', 'query'])
+    with open(DIR + '/sql_in_python_variables_example_correct.py', 'r') as inp:
+        assert inp.read() == example