CLI/git: add @timed decorator, time a few functions

[notmuch] / notmuch-git.py

diff --git a/notmuch-git.py b/notmuch-git.py
index 5475d0dbd26328da380d9ec20994f395c06c6667..6ea50fe8f4ebeef450040bf31c66fca43eed5f7b 100644 (file)
--- a/notmuch-git.py
+++ b/notmuch-git.py
@@ -18,13 +18,6 @@
 
 """
 Manage notmuch tags with Git
-
-Environment variables:
-
-* NMBGIT specifies the location of the git repository used by nmbug.
-  If not specified $HOME/.nmbug is used.
-* NMBPREFIX specifies the prefix in the notmuch database for tags of
-  interest to nmbug. If not specified 'notmuch::' is used.
 """
 
 from __future__ import print_function
@@ -50,7 +43,7 @@ _LOG = _logging.getLogger('nmbug')
 _LOG.setLevel(_logging.WARNING)
 _LOG.addHandler(_logging.StreamHandler())
 
-NMBGIT = None
+NOTMUCH_GIT_DIR = None
 TAG_PREFIX = None
 
 _HEX_ESCAPE_REGEX = _re.compile('%[0-9A-F]{2}')
@@ -96,6 +89,20 @@ def _xapian_unquote(string):
     return string
 
 
+def timed(fn):
+    """Timer decorator"""
+    from time import perf_counter
+
+    def inner(*args, **kwargs):
+        start_time = perf_counter()
+        rval = fn(*args, **kwargs)
+        end_time = perf_counter()
+        _LOG.info('{0}: {1:.8f}s elapsed'.format(fn.__name__, end_time - start_time))
+        return rval
+
+    return inner
+
+
 class SubprocessError(RuntimeError):
     "A subprocess exited with a nonzero status"
     def __init__(self, args, status, stdout=None, stderr=None):
@@ -198,7 +205,7 @@ def _spawn(args, input=None, additional_env=None, wait=False, stdin=None,
 
 
 def _git(args, **kwargs):
-    args = ['git', '--git-dir', NMBGIT] + list(args)
+    args = ['git', '--git-dir', NOTMUCH_GIT_DIR] + list(args)
     return _spawn(args=args, **kwargs)
 
 
@@ -226,16 +233,17 @@ def _get_remote():
         stdout=_subprocess.PIPE, wait=True)
     return remote.strip()
 
+def _tag_query(prefix=None):
+    if prefix is None:
+        prefix = TAG_PREFIX
+    return '(tag (starts-with "{:s}"))'.format(prefix.replace('"','\\\"'))
 
 def get_tags(prefix=None):
     "Get a list of tags with a given prefix."
-    if prefix is None:
-        prefix = TAG_PREFIX
     (status, stdout, stderr) = _spawn(
-        args=['notmuch', 'search', '--output=tags', '*'],
+        args=['notmuch', 'search', '--query=sexp', '--output=tags', _tag_query(prefix)],
         stdout=_subprocess.PIPE, wait=True)
-    return [tag for tag in stdout.splitlines() if tag.startswith(prefix)]
-
+    return [tag for tag in stdout.splitlines()]
 
 def archive(treeish='HEAD', args=()):
     """
@@ -266,7 +274,7 @@ def clone(repository):
     with _tempfile.TemporaryDirectory(prefix='nmbug-clone.') as workdir:
         _spawn(
             args=[
-                'git', 'clone', '--no-checkout', '--separate-git-dir', NMBGIT,
+                'git', 'clone', '--no-checkout', '--separate-git-dir', NOTMUCH_GIT_DIR,
                 repository, workdir],
             wait=True)
     _git(args=['config', '--unset', 'core.worktree'], wait=True, expect=(0, 5))
@@ -327,6 +335,7 @@ def commit(treeish='HEAD', message=None):
         _git(args=['read-tree', treeish], wait=True)
         raise
 
+@timed
 def _update_index(status):
     with _git(
             args=['update-index', '--index-info'],
@@ -358,7 +367,8 @@ def init(remote=None):
     This wraps 'git init' with a few extra steps to support subsequent
     status and commit commands.
     """
-    _spawn(args=['git', '--git-dir', NMBGIT, 'init', '--bare'], wait=True)
+    _spawn(args=['git', '--git-dir', NOTMUCH_GIT_DIR, 'init',
+                 '--initial-branch=master', '--quiet', '--bare'], wait=True)
     _git(args=['config', 'core.logallrefupdates', 'true'], wait=True)
     # create an empty blob (e69de29bb2d1d6434b8b29ae775ad8c2e48c5391)
     _git(args=['hash-object', '-w', '--stdin'], input='', wait=True)
@@ -366,7 +376,7 @@ def init(remote=None):
         args=[
             'commit', '--allow-empty', '-m', 'Start a new nmbug repository'
         ],
-        additional_env={'GIT_WORK_TREE': NMBGIT},
+        additional_env={'GIT_WORK_TREE': NOTMUCH_GIT_DIR},
         wait=True)
 
 
@@ -566,6 +576,7 @@ def _is_unmerged(ref='@{upstream}'):
     return base != fetch_head
 
 
+@timed
 def get_status():
     status = {
         'deleted': {},
@@ -586,17 +597,16 @@ def get_status():
     _os.remove(index)
     return status
 
-
+@timed
 def _index_tags():
     "Write notmuch tags to the nmbug.index."
-    path = _os.path.join(NMBGIT, 'nmbug.index')
-    query = ' '.join('tag:"{tag}"'.format(tag=tag) for tag in get_tags())
+    path = _os.path.join(NOTMUCH_GIT_DIR, 'nmbug.index')
     prefix = '+{0}'.format(_ENCODED_TAG_PREFIX)
     _git(
         args=['read-tree', '--empty'],
         additional_env={'GIT_INDEX_FILE': path}, wait=True)
     with _spawn(
-            args=['notmuch', 'dump', '--format=batch-tag', '--', query],
+            args=['notmuch', 'dump', '--format=batch-tag', '--query=sexp', '--', _tag_query()],
             stdout=_subprocess.PIPE) as notmuch:
         with _git(
                 args=['update-index', '--index-info'],
@@ -636,6 +646,7 @@ def _index_tags_for_message(id, status, tags):
         yield '{mode} {hash}\t{path}\n'.format(mode=mode, hash=hash, path=path)
 
 
+@timed
 def _diff_index(index, filter):
     """
     Get an {id: {tag, ...}} dict for a given filter.
@@ -696,6 +707,14 @@ def _help(parser, command=None):
     else:
         parser.parse_args(['--help'])
 
+def _notmuch_config_get(key):
+    (status, stdout, stderr) = _spawn(
+        args=['notmuch', 'config', 'get', key],
+        stdout=_subprocess.PIPE, wait=True)
+    if status != 0:
+        _LOG.error("failed to run notmuch config")
+        sys.exit(1)
+    return stdout.rstrip()
 
 if __name__ == '__main__':
     import argparse
@@ -708,7 +727,7 @@ if __name__ == '__main__':
         help='Git repository to operate on.')
     parser.add_argument(
         '-p', '--tag-prefix', metavar='PREFIX',
-        default = _os.getenv('NMBPREFIX', 'notmuch::'),
+        default = _os.getenv('NOTMUCH_GIT_PREFIX', 'notmuch::'),
         help='Prefix of tags to operate on.')
     parser.add_argument(
         '-l', '--log-level',
@@ -821,13 +840,13 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     if args.git_dir:
-        NMBGIT = args.git_dir
+        NOTMUCH_GIT_DIR = args.git_dir
     else:
-        NMBGIT = _os.path.expanduser(
-        _os.getenv('NMBGIT', _os.path.join('~', '.nmbug')))
-        _NMBGIT = _os.path.join(NMBGIT, '.git')
-        if _os.path.isdir(_NMBGIT):
-            NMBGIT = _NMBGIT
+        NOTMUCH_GIT_DIR = _os.path.expanduser(
+        _os.getenv('NOTMUCH_GIT_DIR', _os.path.join('~', '.nmbug')))
+        _NOTMUCH_GIT_DIR = _os.path.join(NOTMUCH_GIT_DIR, '.git')
+        if _os.path.isdir(_NOTMUCH_GIT_DIR):
+            NOTMUCH_GIT_DIR = _NOTMUCH_GIT_DIR
 
     TAG_PREFIX = args.tag_prefix
     _ENCODED_TAG_PREFIX = _hex_quote(TAG_PREFIX, safe='+@=,')  # quote ':'
@@ -837,9 +856,13 @@ if __name__ == '__main__':
         _LOG.setLevel(level)
 
     # for test suite
-    for var in ['NMBGIT', 'NMBPREFIX', 'NOTMUCH_PROFILE', 'NOTMUCH_CONFIG' ]:
+    for var in ['NOTMUCH_GIT_DIR', 'NOTMUCH_GIT_PREFIX', 'NOTMUCH_PROFILE', 'NOTMUCH_CONFIG' ]:
         _LOG.debug('env {:s} = {:s}'.format(var, _os.getenv(var,'%None%')))
 
+    if _notmuch_config_get('built_with.sexp_queries') != 'true':
+        _LOG.error("notmuch git needs sexp query support")
+        sys.exit(1)
+
     if not getattr(args, 'func', None):
         parser.print_usage()
         _sys.exit(1)