add method get_filenames_in_dir in utils.py, add unittest

Author: supercoderhawk

pysenal/utils/utils.py

@@ -1,4 +1,5 @@
 # -*- coding: UTF-8 -*-
+import os
 import copy
 from collections import Iterable
 
@@ -48,3 +49,49 @@ def list2dict(l, key, pop_key=False):
             new_item.pop(key)
         d[item[key]] = new_item
     return d
+
+
+def get_filenames_in_dir(dirname,
+                         skip_dir=True,
+                         skip_hidden_file=True,
+                         rm_extname=False,
+                         rm_dirname=False):
+    """
+    get file names in directory, file name in dictionary order
+    :param dirname: directory path to get file names
+    :param skip_dir: whether skip directory in results
+    :param skip_hidden_file: whether skip hidden file in results
+    :param rm_extname: whether remove suffix extname in results
+    :param rm_dirname: whether remove prefix dirname in results
+    :return: scanned file name list, in dictionary order
+    """
+    filenames = []
+    if not os.path.exists(dirname):
+        raise FileNotFoundError('directory is not existed.')
+
+    for filename in os.listdir(dirname):
+        is_add = True
+        if skip_hidden_file and filename.startswith('.') and filename not in {'.', '..'}:
+            is_add = False
+        if skip_dir and os.path.isdir(os.path.join(dirname, filename)):
+            is_add = False
+        if is_add:
+            filenames.append(filename)
+
+    if not rm_dirname:
+        filenames = [os.path.join(dirname, name) for name in filenames]
+
+    if rm_extname:
+        new_filenames = []
+        for filename in filenames:
+            basename = os.path.basename(filename)
+            if '.' in basename:
+                dot_index = basename.rindex('.')
+                if dot_index:
+                    filename = filename[:filename.rindex('.')]
+
+            new_filenames.append(filename)
+        filenames = new_filenames
+
+    filenames = sorted(filenames)
+    return filenames

tests/utils/test_utils.py

@@ -1,4 +1,6 @@
 # -*- coding: UTF-8 -*-
+import shutil
+import tempfile
 from types import GeneratorType
 import pytest
 from pysenal.utils.utils import *
@@ -56,3 +58,40 @@ def test_list2dict():
     assert list2dict(l2, 'name', pop_key=True) == expected_l2
     with pytest.raises(TypeError):
         list2dict(1, 'pid')
+
+
+@pytest.fixture()
+def get_filename_dirname():
+    base_dirname = tempfile.gettempdir() + '/'
+    dirname = base_dirname + 'pysenal_test_utils_get_filename/'
+    return dirname
+
+
+@pytest.fixture()
+def get_filename_dir_setup(get_filename_dirname):
+    dirname = get_filename_dirname
+    if os.path.exists(dirname):
+        shutil.rmtree(dirname)
+    os.mkdir(dirname)
+    os.mkdir(dirname + 'dir1')
+    os.mkdir(dirname + 'dir2')
+    open(dirname + 'file1', 'a').close()
+    open(dirname + 'file2.bak', 'a').close()
+    open(dirname + 'file2.bak', 'a').close()
+    open(dirname + '.file3', 'a').close()
+    open(dirname + '.file4.txt', 'a').close()
+
+
+def test_get_filename_in_dir(get_filename_dir_setup, get_filename_dirname):
+    dirname = get_filename_dirname
+    expected_ret1 = [dirname + 'file1', dirname + 'file2.bak']
+    assert get_filenames_in_dir(dirname) == expected_ret1
+
+    expected_ret2 = [dirname + 'file1', dirname + 'file2']
+    assert get_filenames_in_dir(dirname, rm_extname=True) == expected_ret2
+
+    assert get_filenames_in_dir(dirname, rm_extname=True, rm_dirname=True) == ['file1', 'file2']
+
+    ret3 = get_filenames_in_dir(dirname, rm_extname=True, skip_hidden_file=False)
+    expected_ret3 = [dirname + '.file3', dirname + '.file4', dirname + 'file1', dirname + 'file2']
+    assert ret3 == expected_ret3