Improves the search_file function

Speeds up searching and fixes #88

bob/extension/download.py

@@ -2,6 +2,7 @@
 # vim: set fileencoding=utf-8 :
 
 import bz2
+import glob
 import hashlib
 import io
 import logging
@@ -347,21 +348,22 @@ def find_element_in_tarball(filename, target_path, open_as_stream=False):
     """
 
     f = tarfile.open(filename)
-    for member in f.getmembers():
-        if member.isdir():
+    # iterate over the members of the tarball
+    while True:
+        member = f.next()
+        if member is None:
+            return None
+
+        if not member.isfile():
             continue
 
-        if (
-            member.isfile()
-            and target_path in member.name
-            and os.path.split(target_path)[-1] == os.path.split(member.name)[-1]
-        ):
-            if open_as_stream:
-                return io.BufferedReader(f.extractfile(member)).read()
-            else:
-                return io.TextIOWrapper(f.extractfile(member), encoding="utf-8")
-    else:
-        return None
+        if not member.name.endswith(target_path):
+            continue
+
+        if open_as_stream:
+            return io.BufferedReader(f.extractfile(member)).read()
+        else:
+            return io.TextIOWrapper(f.extractfile(member), encoding="utf-8")
 
 
 def search_file(base_path, options):
@@ -372,16 +374,18 @@ def search_file(base_path, options):
     ----------
 
     base_path: str
-       Base path to start the search, or the tarball to be searched
+        Base folder to start the search, or the tarball to be searched
 
     options: list
-       Files to be searched. This function will return the first occurency
+        Files to be searched. This function will return the first occurrence.
+        The option can be an incomplete relative path. For example, if you have
+        a file called ``"/a/b/c/d.txt"``, and base_path is ``"/a/b"``, then
+        options can be ``["d.txt"]``.
 
     Returns
     -------
     object
         It returns an opened file
-
     """
 
     if not isinstance(options, list):
@@ -389,26 +393,13 @@ def search_file(base_path, options):
 
     # If the input is a directory
     if os.path.isdir(base_path):
-
-        def get_fs():
-            fs = []
-            for root, _, files in os.walk(base_path, topdown=False):
-                for name in files:
-                    fs.append(os.path.join(root, name))
-            return fs
-
-        def search_in_list(o, lst):
-            for i, l in enumerate(lst):
-                if o in l:
-                    return i
-            else:
-                return -1
-
-        fs = get_fs()
         for o in options:
-            index = search_in_list(o, fs)
-            if index >= 0:
-                return open(fs[index])
+            # we append './' to o because o might start with /
+            pattern = os.path.join(base_path, "**", f"./{o}")
+            for path in glob.iglob(pattern, recursive=True):
+                if not os.path.isfile(path):
+                    continue
+                return open(path)
         else:
             return None
     else:

bob/extension/test_download.py

@@ -97,27 +97,37 @@ def test_search_file():
     filename = pkg_resources.resource_filename(
         __name__, "data/example_csv_filelist.tar.gz"
     )
-    # Search in the tarball
-    assert (
-        search_file(filename, "protocol_dev_eval/norm/train_world.csv")
-        is not None
-    )
-    assert search_file(filename, "protocol_dev_eval/norm/xuxa.csv") is None
-
-    # Search in a file structure
-    final_path = "./test_search_file"
-
-    pass
-
-    _untar(filename, final_path, ".gz")
-
-    assert (
-        search_file(final_path, "protocol_dev_eval/norm/train_world.csv")
-        is not None
-    )
-    assert search_file(final_path, "protocol_dev_eval/norm/xuxa.csv") is None
 
-    shutil.rmtree(final_path)
+    with tempfile.TemporaryDirectory(suffix="_extracted") as tmpdir:
+
+        _untar(filename, tmpdir, ".gz")
+
+        # Search in the tarball and in its extracted folder
+        for final_path in (filename, tmpdir):
+            in_extracted_folder = final_path.endswith("_extracted")
+            all_files = list_dir(final_path)
+
+            output_file = search_file(
+                final_path, "protocol_dev_eval/norm/train_world.csv"
+            )
+            assert output_file is not None, all_files
+
+            # test to see if using / we can force an exact match
+            output_file = search_file(
+                final_path, "/protocol_dev_eval/norm/train_world.csv"
+            )
+            assert output_file is not None, all_files
+            assert "my_data" not in output_file.read()
+            if in_extracted_folder:
+                assert "my_protocol" not in output_file.name
+
+            assert (
+                search_file(final_path, "norm/train_world.csv") is not None
+            ), all_files
+            assert (
+                search_file(final_path, "protocol_dev_eval/norm/xuxa.csv")
+                is None
+            ), all_files
 
 
 def test_list_dir():