Optimise string handling in ansible.utils._clean_data

lib/ansible/utils/__init__.py

@@ -33,7 +33,7 @@ from ansible.module_utils.splitter import split_args, unquote
 import ansible.constants as C
 import ast
 import time
-import StringIO
+import cStringIO
 import stat
 import termios
 import tty
@@ -46,6 +46,7 @@ import getpass
 import sys
 import json
 import subprocess
+import contextlib
 
 from vault import VaultLib
 
@@ -55,7 +56,9 @@ MAX_FILE_SIZE_FOR_DIFF=1*1024*1024
 
 # caching the compilation of the regex used
 # to check for lookup calls within data
-LOOKUP_REGEX=re.compile(r'lookup\s*\(')
+LOOKUP_REGEX = re.compile(r'lookup\s*\(')
+PRINT_CODE_REGEX = re.compile(r'(?:{[{%]|[%}]})')
+CODE_REGEX = re.compile(r'(?:{%|%})')
 
 try:
     import json
@@ -355,64 +358,55 @@ def _clean_data(orig_data, from_remote=False, from_inventory=False):
     if not isinstance(orig_data, basestring):
         return orig_data
 
-    data = StringIO.StringIO("")
-
     # when the data is marked as having come from a remote, we always
     # replace any print blocks (ie. {{var}}), however when marked as coming
     # from inventory we only replace print blocks that contain a call to
     # a lookup plugin (ie. {{lookup('foo','bar'))}})
     replace_prints = from_remote or (from_inventory and '{{' in orig_data and LOOKUP_REGEX.search(orig_data) is not None)
 
-    # these variables keep track of opening block locations, as we only
+    regex = PRINT_CODE_REGEX if replace_prints else ONLY_CODE_REGEX
-    # want to replace matched pairs of print/block tags
+
-    print_openings = []
+    with contextlib.closing(cStringIO.StringIO()) as data:
-    block_openings = []
+        # these variables keep track of opening block locations, as we only
-
+        # want to replace matched pairs of print/block tags
-    for idx,c in enumerate(orig_data):
+        last_pos = 0
-        # if the current character is an opening brace, check to
+        print_openings = []
-        # see if this is a jinja2 token. Otherwise, if the current
+        block_openings = []
-        # character is a closing brace, we backup one character to
+        for mo in regex.finditer(orig_data):
-        # see if we have a closing.
+            token = mo.group(0)
-        if c == '{' and idx < len(orig_data) - 1:
+            token_start = mo.start(0)
-            token = orig_data[idx:idx+2]
+            token_end = mo.end(0)
-            # if so, and we want to replace this block, push
+
-            # this token's location onto the appropriate array
+            if token[0] == '{':
-            if token == '{{' and replace_prints:
+                if token == '{%':
-                print_openings.append(idx)
+                    block_openings.append(token_start)
-            elif token == '{%':
+                elif token == '{{':
-                block_openings.append(idx)
+                    print_openings.append(token_start)
-            # finally we write the data to the buffer and write
+                data.write(orig_data[last_pos:token_end])
-            data.seek(0, os.SEEK_END)
+            elif token[1] == '}':
-            data.write(c)
+                prev_idx = None
-        elif c == '}' and idx > 0:
+                if token == '%}' and block_openings:
-            token = orig_data[idx-1:idx+1]
+                    prev_idx = block_openings.pop()
-            prev_idx = -1
+                elif token == '}}' and print_openings:
-            if token == '%}' and len(block_openings) > 0:
+                    prev_idx = print_openings.pop()
-                prev_idx = block_openings.pop()
+
-            elif token == '}}' and len(print_openings) > 0:
+                data.write(orig_data[last_pos:token_start])
-                prev_idx = print_openings.pop()
+                if prev_idx is not None:
-            # if we have a closing token, and we have previously found
+                    # replace the opening
-            # the opening to the same kind of block represented by this
+                    data.seek(prev_idx, os.SEEK_SET)
-            # token, replace both occurrences, otherwise we just write
+                    data.write('{#')
-            # the current character to the buffer
+                    # replace the closing
-            if prev_idx != -1:
+                    data.seek(0, os.SEEK_END)
-                # replace the opening
+                    data.write('#}')
-                data.seek(prev_idx, os.SEEK_SET)
+                else:
-                data.write('{#')
+                    data.write(token)
-                # replace the closing
-                data.seek(-1, os.SEEK_END)
-                data.write('#}')
             else:
-                data.seek(0, os.SEEK_END)
+                assert False, 'Unhandled regex match'
-                data.write(c)
+            last_pos = token_end
-        else:
+
-            # not a jinja2 token, so we just write the current char
+        data.write(orig_data[last_pos:])
-            # to the output buffer
+
-            data.seek(0, os.SEEK_END)
+        return data.getvalue()
-            data.write(c)
-    return_data = data.getvalue()
-    data.close()
-    return return_data
 
 def _clean_data_struct(orig_data, from_remote=False, from_inventory=False):
     '''