Merge pull request #3 from Smirl/processes

Processes

README.md

@@ -3,9 +3,7 @@
 Toolkit for easy searching and manipulation of python source code using
 redbaron.
 
-[ ![Codeship Status for Smirl/baroness](https://app.codeship.com/projects/647edcd0-bcc0-0135-7fde-5afd35787ded/status?branch=master)](https://app.codeship.com/projects/259596)
-
-[![Coverage Status](https://coveralls.io/repos/github/Smirl/baroness/badge.svg?branch=HEAD)](https://coveralls.io/github/Smirl/baroness?branch=HEAD)
+[ ![Codeship Status for Smirl/baroness](https://app.codeship.com/projects/647edcd0-bcc0-0135-7fde-5afd35787ded/status?branch=master)](https://app.codeship.com/projects/259596) [![Coverage Status](https://coveralls.io/repos/github/Smirl/baroness/badge.svg?branch=HEAD)](https://coveralls.io/github/Smirl/baroness?branch=HEAD)
 
 ## Installation
 

baroness/cli.py

@@ -3,10 +3,12 @@
 
 from __future__ import print_function
 from argparse import ArgumentParser
+import multiprocessing
+import logging
 
 from baroness.cache import cache_save, cache_delete
 from baroness.search import search
-from baroness.utils import set_default_subparser
+from baroness.utils import set_default_subparser, LOGGER
 
 
 def main():
@@ -74,10 +76,24 @@ def main():
         action='store_true',
         help='Do not output the linenumbers.'
     )
+    search_parser.add_argument(
+        '-v', '--verbose',
+        default=0,
+        action='count',
+        help='Level of verbosity. Can be used many times.'
+    )
+    search_parser.add_argument(
+        '-P', '--processes',
+        default=multiprocessing.cpu_count() + 1,  # 'cus why not
+        type=int,
+        help='Number of processes to use.'
+    )
     search_parser.set_defaults(func=search)
 
     parser.set_default_subparser('search')
 
     args = vars(parser.parse_args())
     func = args.pop('func')
+    if args.get('verbose', 0) > 0:
+        LOGGER.setLevel(logging.DEBUG)
     return func(**args)

baroness/search.py

@@ -2,9 +2,11 @@
 
 
 from __future__ import print_function
+from multiprocessing import Pool
 import os
+import signal
+import sys
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from redbaron import RedBaron
 from redbaron.base_nodes import NodeList
 
@@ -19,46 +21,85 @@
     import json
 
 
-def search(pattern, files, no_cache, parents, no_color, no_linenos):
+def search(
+    files, pattern, no_cache=False, parents=0, no_color=False,
+    no_linenos=False, verbose=0, processes=2
+):
     """Search all files with the redbaron expression in pattern."""
-    local = locals()
-    exec('def _search(root):\n    return {}'.format(pattern), globals(), local)
-    _search = local['_search']
+    options = dict(locals())
+    options.pop('files')
 
-    def _search_file(filename):
+    LOGGER.debug('DEBUG Using options: %s', options)
+    LOGGER.debug('DEBUG Using %s processes', processes)
+    pool = Pool(processes, _init_worker)
+    tasks = ((filename, pattern, options) for filename in filenames(files))
+    try:
+        for result in pool.imap_unordered(_safe_search_file, tasks):
+            if not result:
+                continue
+            filename, output = result
+            LOGGER.debug('DEBUG Searching: %s', filename)
+            if output:
+                LOGGER.info(output)
+    except KeyboardInterrupt:
+        LOGGER.critical('KeyboardInterrupt, quitting')
+        sys.exit(1)
+    finally:
+        pool.terminate()
+        pool.join()
+
+
+def _init_worker():
+    """Ensure that KeyboardInterrupt is ignored."""
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+
+def _safe_search_file(args):
+    """A wrapper for multiprocessing around _search_file."""
+    try:
+        return _search_file(*args)
+    except Exception:
+        LOGGER.exception('Error processing file')
+
+
+def _search_file(filename, pattern, options):
         """Search a given filename."""
+        local = locals()
+        exec('def _search(root):\n    return {}'.format(pattern), globals(), local)
+        _search = local['_search']
+
         cache_file = _cache_filename(filename)
-        if not no_cache and os.path.exists(cache_file):
+        if not options['no_cache'] and os.path.exists(cache_file):
             with open(cache_file) as f:
                 root = RedBaron(NodeList.from_fst(json.load(f)))
             for node in root:
                 node.parent = root
             root.node_list.parent = root
         else:
-            root = _load_and_save(filename, cache_file, no_cache=no_cache)
+            root = _load_and_save(filename, cache_file, no_cache=options['no_cache'])
 
         results = _search(root)
         output = []
+        seen = set()
         if results:
             output.append(filename)
             for result in results:
-                for _ in range(parents):
+                # Get the correct number of parents
+                for _ in range(options['parents']):
                     result = result.parent if result.parent else result
-                output.append(format_node(result, no_color=no_color, no_linenos=no_linenos))
-                output.append('--')
-            output.append('')
-        return output
-
-    with ThreadPoolExecutor(max_workers=10) as executor:
-        tasks = [
-            executor.submit(_search_file, filename)
-            for filename in filenames(files)
-        ]
-        for future in as_completed(tasks):
-            try:
-                output = '\n'.join(future.result())
-            except Exception:
-                LOGGER.exception('Error processing file')
-            else:
-                if output:
-                    LOGGER.info(output)
+
+                # Ensure that we don't print the same code twice
+                if result in seen:
+                    continue
+                else:
+                    seen.add(result)
+
+                # format the output
+                output.append(format_node(
+                    result,
+                    no_color=options['no_color'],
+                    no_linenos=options['no_linenos']
+                ))
+                output.append(u'--')
+            output.append(u'')
+        return (filename, u'\n'.join(output))

baroness/utils.py

@@ -19,7 +19,7 @@
 
 
 LOGGER = logging.getLogger('baroness')
-LOGGER.setLevel(logging.DEBUG)
+LOGGER.setLevel(logging.INFO)
 ch = logging.StreamHandler()
 ch.setLevel(logging.DEBUG)
 LOGGER.addHandler(ch)
@@ -33,6 +33,7 @@ def filenames(files):
     and yield python files. Or can be a list of filenames and/or glob patterns
     which we expand, and filter for python files.
     """
+    LOGGER.debug('Expanding globs: %s', files)
     directories = []
     if files:
         for pattern in files:
@@ -45,6 +46,7 @@ def filenames(files):
     if not (files or directories):
         directories.append('.')
 
+    LOGGER.debug('Searching for python files in: %s', directories)
     for directory in directories:
         for root, dirs, files in os.walk(directory):
             for filename in files:
@@ -95,6 +97,8 @@ def format_lineno(lineno):
             linenos = count(node.absolute_bounding_box.top_left.line)
         except AttributeError:
             linenos = repeat(u'**')
+        except ValueError:
+            linenos = count()
         lines = [
             u'{}-{}'.format(format_lineno(lineno).strip(), line)
             for lineno, line in zip(linenos, lines)
@@ -118,4 +122,9 @@ def _load_and_save(filename, cache_file, no_cache=False):
 
 def _cache_filename(filename):
     """Return the filename to store the cached data in."""
-    return os.path.join('.baroness', filename) + '.json'
+    head, tail = os.path.split(filename)
+    return os.path.join(
+        '.baroness',
+        os.path.abspath(head).lstrip(os.path.sep),
+        tail + '.json'
+    )

tests/baroness/search_test.py

@@ -0,0 +1,32 @@
+"""Test the search functionality."""
+
+import logging
+
+import pytest
+
+from baroness.search import search
+from baroness.utils import LOGGER
+
+LOGGER.setLevel(logging.DEBUG)
+
+
+@pytest.fixture
+def workdir(tmpdir):
+    """
+    A place to add files to search.
+
+    Create a directory called package/
+    """
+    tmpdir.mkdir('package')
+    return tmpdir
+
+
+@pytest.fixture
+def python_file(workdir):
+    """Set up temporary files and return the tmpdir."""
+    workdir.join('main.py').write('''print("hello")\n''')
+
+
+def test_search(workdir, python_file, capsys):
+    """Test the standard code path which yeilds results."""
+    search([str(workdir)], 'root("name")', verbose=1)