Improve performance on pre-commit -r pyflakes by 40%

Anthony Sottile · Anthony Sottile · commit da8c435e9c89 · 2014-03-13T23:48:17.000-07:00
diff --git a/pre_commit/git.py b/pre_commit/git.py
@@ -1,14 +1,16 @@
 import os
 import pkg_resources
-
 from plumbum import local
 
+from pre_commit.util import memoize_by_cwd
+
 
-# TODO: optimization: memoize based on local.cwd.getpath()
+@memoize_by_cwd
 def get_root():
     return local['git']['rev-parse', '--show-toplevel']().strip()
 
 
+@memoize_by_cwd
 def get_pre_commit_path():
     return os.path.join(get_root(), '.git/hooks/pre-commit')
 
@@ -25,4 +27,4 @@ def remove_pre_commit():
 
 def get_head_sha(git_repo_path):
     with local.cwd(git_repo_path):
-        return (local['git']['rev-parse', 'HEAD'])().strip()
+        return (local['git']['rev-parse', 'HEAD'])().strip()
diff --git a/pre_commit/util.py b/pre_commit/util.py
@@ -1,4 +1,8 @@
 
+import functools
+import os
+
+
 class cached_property(object):
     """Like @property, but caches the value."""
 
@@ -14,3 +18,19 @@ def __get__(self, obj, cls):
         value = self._func(obj)
         obj.__dict__[self.__name__] = value
         return value
+
+
+def memoize_by_cwd(func):
+    """Memoize a function call based on os.getcwd()."""
+    cache = {}
+    @functools.wraps(func)
+    def wrapper(*args):
+        cwd = os.getcwd()
+        key = (cwd,) + args
+        try:
+            return cache[key]
+        except KeyError:
+            ret = cache[key] = func(*args)
+            return ret
+
+    return wrapper
diff --git a/tests/util_test.py b/tests/util_test.py
@@ -1,16 +1,18 @@
 
 import pytest
-import time
+import random
+from plumbum import local
 
 from pre_commit.util import cached_property
+from pre_commit.util import memoize_by_cwd
 
 
 @pytest.fixture
 def class_with_cached_property():
     class Foo(object):
         @cached_property
         def foo(self):
-            return "Foo" + str(time.time())
+            return "Foo" + str(random.getrandbits(64))
 
     return Foo
 
@@ -27,3 +29,33 @@ def test_unbound_cached_property(class_with_cached_property):
     prop = class_with_cached_property.foo
     assert isinstance(prop, cached_property)
 
+
+@pytest.fixture
+def memoized_by_cwd():
+    @memoize_by_cwd
+    def func(arg):
+        return arg + str(random.getrandbits(64))
+
+    return func
+
+
+def test_memoized_by_cwd_returns_same_twice_in_a_row(memoized_by_cwd):
+    ret = memoized_by_cwd('baz')
+    ret2 = memoized_by_cwd('baz')
+    assert ret is ret2
+
+
+def test_memoized_by_cwd_returns_different_for_different_args(memoized_by_cwd):
+    ret = memoized_by_cwd('baz')
+    ret2 = memoized_by_cwd('bar')
+    assert ret.startswith('baz')
+    assert ret2.startswith('bar')
+    assert ret != ret2
+
+
+def test_memoized_by_cwd_changes_with_different_cwd(memoized_by_cwd):
+    ret = memoized_by_cwd('baz')
+    with local.cwd('.git'):
+        ret2 = memoized_by_cwd('baz')
+
+    assert ret != ret2
