Complete TODO: handle custom modification functions as well as regexp

src-d#526 (comment)

Signed-off-by: egor <[email protected]>

lookout/style/format/benchmarks/generate_smoke.py

@@ -6,6 +6,7 @@
 import re
 import shutil
 import tarfile
+from types import FunctionType
 from typing import Tuple
 
 from dulwich.index import blob_from_path_and_stat, build_index_from_tree
@@ -103,14 +104,20 @@ def generate_smoke_entry(inputpath: str, outputpath: str, force: bool = False) -
         writer = csv.DictWriter(index_file, fieldnames=["repo", "style", "from", "to"])
         writer.writeheader()
         for style_name, (pattern, repl) in js_format_rules.items():
-            # TODO(zurk): handle custom modification functions as well as regexp.
-            pattern = re.compile(pattern)
+            if isinstance(pattern, str):
+                pattern_ = re.compile(pattern)
+
+                def pattern(repl, code):
+                    return re.sub(pattern_, repl, code)
+            else:
+                assert isinstance(pattern, FunctionType), "Expected function, got %s" % \
+                                                          type(pattern)
             for repopath in repopaths:
                 repo = Repo(str(repopath))
                 repo.hooks.clear()  # Speed up dulwich by ~25%
                 for filepath in repopath.glob("**/*.js"):
                     code = filepath.read_text()
-                    new_code = re.sub(pattern, repl, code)
+                    new_code = pattern(repl, code)
                     filepath.write_text(new_code)
                 from_commit, to_commit = commit_style(repo, style_name)
                 writer.writerow({