python（遍历）读取文件或文件夹

超、凢脫俗 2022-12-10 04:52 237阅读 0赞

### 文章目录 ###

*  例子
 *  os.listdir
 *  os.walk
 *  遍历读取代码

搞机器学习或者深度学习算法很多时候需要遍历某个目录读取文件，特别是经常需要读取某个特定后缀的文件，比如图片的话可能需要读取`jpg, png, bmp`格式的文件。python本身的库函数功能没有这么定制化，所以就需要再重新包装一下。

# 例子 #

假设我们有如下的目录结构，以`bmp`结尾的是`文件`，其他是`文件夹`。下面的程序都将以该目录结构为例进行说明。

data
    	- a
    		- a1.bmp
    	- b
    		- b1.bmp
    	- 1.bmp
    	- 2.bmp

# os.listdir #

`os.listdir`仅读取当前路径下的文件和文件夹，返回一个列表。读取demo目录结构的代码和结果如下：

path = r'D:\data'
    items = os.listdir(path)  # ==> ['1.bmp', '2.bmp', 'a', 'b']

# os.walk #

`os.walk`本身已经是遍历读取，包含所有的子文件（夹）但是其结果不像是os.listdir一样是个list，而是一个比较复杂的数据体，难以直接使用，所以一般需要再处理一下。我们可以使用for语句将其打印出来看看：

path = r'D:\data'
    # part 1
    for items in os.walk(path):
        print(items)
    # part 2
    for main_dir, sub_dir_list, sub_file_list in os.walk(path):
        print(main_dir, sub_dir_list, sub_file_list)

结果为：

# part 1
    ('D:\\data', ['a', 'b'], ['1.bmp', '2.bmp'])
    ('D:\\data\\a', [], ['a1.bmp'])
    ('D:\\data\\b', [], ['b1.bmp'])
    
    # part 2
    D:\data ['a', 'b'] ['1.bmp', '2.bmp']
    D:\data\a [] ['a1.bmp']
    D:\data\b [] ['b1.bmp']

使用迭代器对`os.walk()`的结果进行输出，发现每一条包含三个部分（part 1），在part 2中，我们给三个部分分别起名为`main_dir, sub_dir_list, sub_file_list`，下面对其进行简单解释：

*  main\_dir：遍历得到的路径下所有文件夹
 *  sub\_dir\_list：main\_dir下面的文件夹
 *  sub\_file\_list：main\_dir下面的文件

连接`main_dir`和`sub_file_list`中的文件可以得到路径下的所有文件。

`sub_dir_list`在这里则没有用处，我们无需再去遍历sub\_dir\_list，因为它们已经包含在main\_dir里了。

# 遍历读取代码 #

代码逻辑如下：

*  需要有后缀辨别功能，并且能够同时辨别多个后缀
 *  需要有递归和非递归功能
 *  返回的是以入参`path`为前缀的路径，所以如果path是完整路径那么返回的就是完整路径，否则就不是

# -*- coding: utf-8 -*-
    import os
    
    
    def file_ext(filename, level=1):
        """ return extension of filename Parameters: ----------- filename: str name of file, path can be included level: int level of extension. for example, if filename is 'sky.png.bak', the 1st level extension is 'bak', and the 2nd level extension is 'png' Returns: -------- extension of filename """
        return filename.split('.')[-level]
    
    
    def _contain_file(path, extensions):
        """ check whether path contains any file whose extension is in extensions list Parameters: ----------- path: str path to be checked extensions: str or list/tuple of str extension or extensions list Returns: -------- return True if contains, else return False """
        assert os.path.exists(path), 'path must exist'
        assert os.path.isdir(path), 'path must be dir'
    
        if isinstance(extensions, str):
            extensions = [extensions]
    
        for file in os.listdir(path):
            if os.path.isfile(os.path.join(path, file)):
                if (extensions is None) or (file_ext(file) in extensions):
                    return True
        return False
    
    
    def _process_extensions(extensions=None):
        """ preprocess and check extensions, if extensions is str, convert it to list. Parameters: ----------- extensions: str or list/tuple of str file extensions Returns: -------- extensions: list/tuple of str file extensions """
        if extensions is not None:
            if isinstance(extensions, str):
                extensions = [extensions]
            assert isinstance(extensions, (list, tuple)), \
                'extensions must be str or list/tuple of str'
            for ext in extensions:
                assert isinstance(ext, str), 'extension must be str'
        return extensions
    
    
    def get_files(path, extensions=None, is_recursive=True):
        """ read files in path. if extensions is None, read all files, if extensions are specified, only read the files who have one of the extensions. if is_recursive is True, recursively read all files, if is_recursive is False, only read files in current path. Parameters: ----------- path: str path to be read extensions: str or list/tuple of str file extensions is_recursive: bool whether read files recursively. read recursively is True, while just read files in current path if False Returns: -------- files: the obtained files in path """
        extensions = _process_extensions(extensions)
        files = []
        # get files in current path
        if not is_recursive:
            for name in os.listdir(path):
                fullname = os.path.join(path, name)
                if os.path.isfile(fullname):
                    if (extensions is None) or (file_ext(fullname) in extensions):
                        files.append(fullname)
            return files
        # get files recursively
        for main_dir, _, sub_file_list in os.walk(path):
            for filename in sub_file_list:
                fullname = os.path.join(main_dir, filename)
                if (extensions is None) or (file_ext(fullname) in extensions):
                    files.append(fullname)
        return files
    
    
    def get_folders(path, extensions=None, is_recursive=True):
        """ read folders in path. if extensions is None, read all folders, if extensions are specified, only read the folders who contain any files that have one of the extensions. if is_recursive is True, recursively read all folders, if is_recursive is False, only read folders in current path. Parameters: ----------- path: str path to be read extensions: str or list/tuple of str file extensions is_recursive: bool whether read folders recursively. read recursively is True, while just read folders in current path if False Returns: -------- folders: the obtained folders in path """
        extensions = _process_extensions(extensions)
        folders = []
        # get folders in current path
        if not is_recursive:
            for name in os.listdir(path):
                fullname = os.path.join(path, name)
                if os.path.isdir(fullname):
                    if (extensions is None) or \
                            (_contain_file(fullname, extensions)):
                        folders.append(fullname)
            return folders
        # get folders recursively
        for main_dir, _, _ in os.walk(path):
            if (extensions is None) or (_contain_file(main_dir, extensions)):
                folders.append(main_dir)
        return folders
    
    
    if __name__ == '__main__':
        path = r'.\data'
    
        files = get_files(path)
        print(files)  # ==> ['D:\\data\\1.bmp', 'D:\\data\\2.bmp', 'D:\\data\\a\\a1.bmp', 'D:\\data\\b\\b1.bmp']
    
        folders = get_folders(path)
        print(folders)  # ==> ['D:\\data', 'D:\\data\\a', 'D:\\data\\b']