modify scandir and add tests

cvbase/io.py

@@ -48,7 +48,7 @@ def pickle_load(filename, **kwargs):
     return obj
 
 
-def pickle_dump(obj, filename, protocol=2, **kwargs):
+def pickle_dump(obj, filename, **kwargs):
     kwargs.setdefault('protocol', 2)
     with open(filename, 'wb') as f:
         pickle.dump(obj, f, **kwargs)
@@ -98,33 +98,31 @@ def mkdir_or_exist(dir_name):
         os.makedirs(dir_name)
 
 
-def _scandir_py35(dir_path, ext=None):
-    if isinstance(ext, str):
-        ext = [ext]
+def _scandir_py35(dir_path, suffix=None):
     for entry in os.scandir(dir_path):
         if not entry.is_file():
             continue
         filename = entry.name
-        if ext is None:
+        if suffix is None:
             yield filename
-        elif filename.split('.')[-1] in ext:
+        elif filename.endswith(suffix):
             yield filename
 
 
-def _scandir_py(dir_path, ext=None):
-    if isinstance(ext, str):
-        ext = [ext]
+def _scandir_py(dir_path, suffix=None):
     for filename in os.listdir(dir_path):
         if not path.isfile(path.join(dir_path, filename)):
             continue
-        if ext is None:
+        if suffix is None:
             yield filename
-        elif filename.split('.')[-1] in ext:
+        elif filename.endswith(suffix):
             yield filename
 
 
-def scandir(dir_path, ext=None):
+def scandir(dir_path, suffix=None):
+    if suffix is not None and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
     if sys.version[0] == 3 and sys.version[1] >= 5:
-        return _scandir_py35(dir_path, ext)
+        return _scandir_py35(dir_path, suffix)
     else:
-        return _scandir_py(dir_path, ext)
+        return _scandir_py(dir_path, suffix)

cvbase/tests/test_io.py

@@ -1,7 +1,7 @@
 from os import makedirs, path, remove, removedirs
 
-from cvbase import (json_dump, json_load, yaml_load, yaml_dump, pickle_dump,
-                    pickle_load, list_from_file, AsyncDumper)
+from cvbase.io import (json_dump, json_load, yaml_load, yaml_dump, pickle_dump,
+                       pickle_load, list_from_file, AsyncDumper, scandir)
 
 
 def test_json():
@@ -71,3 +71,13 @@ def test_async_dumper():
         assert path.isfile(filename)
         remove(filename)
     removedirs(tmp_dir)
+
+
+def test_scandir():
+    folder = path.join(path.dirname(__file__), 'data')
+    assert set(scandir(folder)) == set(
+        ['filelist.txt', 'test.jpg', 'voc_labels.txt'])
+    assert set(scandir(folder, '.txt')) == set(
+        ['filelist.txt', 'voc_labels.txt'])
+    assert set(scandir(folder, ('.jpg', '.txt'))) == set(
+        ['filelist.txt', 'test.jpg', 'voc_labels.txt'])

setup.py

@@ -15,7 +15,7 @@ def readme():
 
 setup(
     name='cvbase',
-    version='0.2.1',
+    version='0.2.2',
     description='Utils for computer vision research',
     long_description=readme(),
     keywords='computer vision',