Refactor hg and git file finders to use the same algorithm

Also make the test parametric.

6 files changed, 98 insertions, 38 deletions

diff --git a/setup.py b/setup.py
index 2b93690..d2cb0f3 100644
--- a/setup.py
+++ b/setup.py
@@ -81,8 +81,8 @@ arguments = dict(
         pip-egg-info = setuptools_scm.hacks:parse_pip_egg_info
 
         [setuptools_scm.files_command]
-        .hg = setuptools_scm.hg:FILES_COMMAND
-        .git = setuptools_scm.git_file_finder:find_files
+        .hg = setuptools_scm.file_finder_hg:hg_find_files
+        .git = setuptools_scm.file_finder_git:git_find_files
 
         [setuptools_scm.version_scheme]
         guess-next-dev = setuptools_scm.version:guess_next_dev_version
diff --git a/setuptools_scm/file_finder.py b/setuptools_scm/file_finder.py
new file mode 100644
index 0000000..fb7648a
--- /dev/null
+++ b/setuptools_scm/file_finder.py
@@ -0,0 +1,33 @@
+import os
+
+
+def scm_find_files(path, scm_files, scm_dirs):
+    """ setuptools compatible file finder that follows symlinks
+
+    - path: the root directory from which to search
+    - scm_files: set of scm controlled files
+    - scm_files: set of scm controlled directories
+
+    scm_files and scm_dirs must be absolute with symlinks resolved (realpath),
+    with normalized case (normcase)
+
+    Spec here: http://setuptools.readthedocs.io/en/latest/setuptools.html#\
+        adding-support-for-revision-control-systems
+    """
+    realpath = os.path.normcase(os.path.realpath(path))
+    seen = set()
+    res = []
+    for dirpath, dirnames, filenames in os.walk(realpath, followlinks=True):
+        # dirpath with symlinks resolved
+        realdirpath = os.path.normcase(os.path.realpath(dirpath))
+        if realdirpath not in scm_dirs or realdirpath in seen:
+            dirnames[:] = []
+            continue
+        for filename in filenames:
+            # dirpath + filename with symlinks preserved
+            fullfilename = os.path.join(dirpath, filename)
+            if os.path.normcase(os.path.realpath(fullfilename)) in scm_files:
+                res.append(
+                    os.path.join(path, os.path.relpath(fullfilename, path)))
+        seen.add(realdirpath)
+    return res
diff --git a/setuptools_scm/git_file_finder.py b/setuptools_scm/file_finder_git.py
index da3d8e7..ecb5cf6 100644
--- a/setuptools_scm/git_file_finder.py
+++ b/setuptools_scm/file_finder_git.py
@@ -2,6 +2,8 @@ import os
 import subprocess
 import tarfile
 
+from .file_finder import scm_find_files
+
 
 def _git_toplevel(path):
     try:
@@ -35,31 +37,9 @@ def _git_ls_files_and_dirs(toplevel):
     return git_files, git_dirs
 
 
-def find_files(path=''):
-    """ setuptools compatible git file finder that follows symlinks
-
-    Spec here: http://setuptools.readthedocs.io/en/latest/setuptools.html#\
-        adding-support-for-revision-control-systems
-    """
+def git_find_files(path=''):
     toplevel = _git_toplevel(path)
     if not toplevel:
         return []
     git_files, git_dirs = _git_ls_files_and_dirs(toplevel)
-    realpath = os.path.normcase(os.path.realpath(path))
-    assert realpath.startswith(toplevel)
-    seen = set()
-    res = []
-    for dirpath, dirnames, filenames in os.walk(realpath, followlinks=True):
-        # dirpath with symlinks resolved
-        realdirpath = os.path.normcase(os.path.realpath(dirpath))
-        if realdirpath not in git_dirs or realdirpath in seen:
-            dirnames[:] = []
-            continue
-        for filename in filenames:
-            # dirpath + filename with symlinks preserved
-            fullfilename = os.path.join(dirpath, filename)
-            if os.path.normcase(os.path.realpath(fullfilename)) in git_files:
-                res.append(
-                    os.path.join(path, os.path.relpath(fullfilename, path)))
-        seen.add(realdirpath)
-    return res
+    return scm_find_files(path, git_files, git_dirs)
diff --git a/setuptools_scm/file_finder_hg.py b/setuptools_scm/file_finder_hg.py
new file mode 100644
index 0000000..2f5ccda
--- /dev/null
+++ b/setuptools_scm/file_finder_hg.py
@@ -0,0 +1,44 @@
+import os
+import subprocess
+
+from .file_finder import scm_find_files
+
+
+def _hg_toplevel(path):
+    try:
+        with open(os.devnull, 'wb') as devnull:
+            out = subprocess.check_output([
+                'hg', 'root',
+            ], cwd=(path or '.'), universal_newlines=True, stderr=devnull)
+        return os.path.normcase(os.path.realpath(out.strip()))
+    except subprocess.CalledProcessError:
+        # hg returned error, we are not in a mercurial repo
+        return None
+    except OSError:
+        # hg command not found, probably
+        return None
+
+
+def _hg_ls_files_and_dirs(toplevel):
+    hg_files = set()
+    hg_dirs = set([toplevel])
+    out = subprocess.check_output([
+        'hg', 'files',
+    ], cwd=toplevel, universal_newlines=True)
+    for name in out.split():
+        name = os.path.normcase(name).replace('/', os.path.sep)
+        fullname = os.path.join(toplevel, name)
+        hg_files.add(fullname)
+        dirname = os.path.dirname(fullname)
+        while len(dirname) > len(toplevel) and dirname not in hg_dirs:
+            hg_dirs.add(dirname)
+            dirname = os.path.dirname(dirname)
+    return hg_files, hg_dirs
+
+
+def hg_find_files(path=''):
+    toplevel = _hg_toplevel(path)
+    if not toplevel:
+        return []
+    hg_files, hg_dirs = _hg_ls_files_and_dirs(toplevel)
+    return scm_find_files(path, hg_files, hg_dirs)
diff --git a/setuptools_scm/hg.py b/setuptools_scm/hg.py
index 422e468..ccf6603 100644
--- a/setuptools_scm/hg.py
+++ b/setuptools_scm/hg.py
@@ -2,8 +2,6 @@ import os
 from .utils import do, trace, data_from_mime, has_command
 from .version import meta, tags_to_versions
 
-FILES_COMMAND = 'hg locate -I .'
-
 
 def _hg_tagdist_normalize_tagcommit(root, tag, dist, node, branch):
     dirty = node.endswith('+')
diff --git a/testing/test_git_file_finder.py b/testing/test_file_finder.py
index fb7d6f1..81a599f 100644
--- a/testing/test_git_file_finder.py
+++ b/testing/test_file_finder.py
@@ -3,16 +3,21 @@ import sys
 
 import pytest
 
-from setuptools_scm.git_file_finder import find_files
-
-
-@pytest.fixture
-def inwd(wd):
-    wd('git init')
-    wd('git config user.email test@example.com')
-    wd('git config user.name "a test"')
-    wd.add_command = 'git add .'
-    wd.commit_command = 'git commit -m test-{reason}'
+from setuptools_scm.integration import find_files
+
+
+@pytest.fixture(params=['git', 'hg'])
+def inwd(request, wd):
+    if request.param == 'git':
+        wd('git init')
+        wd('git config user.email test@example.com')
+        wd('git config user.name "a test"')
+        wd.add_command = 'git add .'
+        wd.commit_command = 'git commit -m test-{reason}'
+    elif request.param == 'hg':
+        wd('hg init')
+        wd.add_command = 'hg add .'
+        wd.commit_command = 'hg commit -m test-{reason} -u test -d "0 0"'
     (wd.cwd / 'file1').ensure(file=True)
     adir = (wd.cwd / 'adir').ensure(dir=True)
     (adir / 'filea').ensure(file=True)