Python：在列表中查找具有匹配扩展名或名称的文件

Question

假设我有一个文件名的列表：[exia.gundam, dynames.gundam, kyrios.gundam, virtue.gundam]，或者[exia.frame, exia.head, exia.swords, exia.legs, exia.arms, exia.pilot, exia.gn_drive, lockon_stratos.data, tieria_erde.data, ribbons_almark.data, otherstuff.dada]。

在一次循环中，我想要获取所有的 *.gundam 或 *.data 文件，而在另一次循环中，我想把所有以 exia. 开头的文件分在一起。除了逐个遍历列表并把每个元素放到字典里，还有什么简单的方法吗？

这是我想到的：

def matching_names(files):
    '''
    extracts files with repeated names from a list

    Keyword arguments:
    files - list of filenames

    Returns: Dictionary
    '''

    nameDict = {}
    for file in files:
        filename = file.partition('.')
        if filename[0] not in nameDict:
            nameDict[filename[0]] = []
        nameDict[filename[0]].append(filename[2])

    matchingDict = {}
    for key in nameDict.keys():
        if len(nameDict[key]) > 1:
            matchingDict[key] = nameDict[key] 
    return matchingDict

那么，假设我必须使用这个，有没有简单的方法可以把文件扩展名作为键，而不是文件名呢？

Answer 1

假设你想要的结果是一个文件名的列表，每个列表里包含了一组文件名，这些文件名是根据扩展名或者根名称来分组的：

import os.path
import itertools as it

def files_grouped_by(filenames, use_extension=True):
    def ky(fn): return os.path.splitext(fn)[use_extension]
    return [list(g) for _, g in it.groupby(sorted(filenames, key=ky), ky)]

现在，调用 files_grouped_by(filenames, False) 会返回一个根据根名称分组的列表，而如果第二个参数是 True 或者没有提供，那么分组就是根据扩展名来进行的。

如果你想要的是一个字典，字典的键是根名称或者扩展名，值则是对应的文件名列表，做法也很相似：

import os.path
import itertools as it

def dict_files_grouped_by(filenames, use_extension=True):
    def ky(fn): return os.path.splitext(fn)[use_extension]
    return dict((k, list(g)) 
                for k, g in it.groupby(sorted(filenames, key=ky), ky)]

Answer 2

我不太确定你想要做什么，但如果我理解得没错，像这样可能会有效：

from collections import defaultdict
files_by_extension = defaultdict(list)

for f in files:
    files_by_extension[ f.split('.')[1] ].append(f)

这个代码是根据文件的扩展名创建一个哈希表，并通过一次遍历列表来填充这个哈希表。

Answer 3

在我第一版的回答中，看来我误解了你的问题。所以如果我理解正确的话，你是想处理一堆文件，这样你就可以轻松找到所有带有特定后缀的文件名，或者所有带有特定基础名的文件名（“基础名”就是指在点之前的部分）？

如果是这样的话，我建议你可以这样做：

from itertools import groupby

def group_by_name(filenames):
    '''Puts the filenames in the given iterable into a dictionary where
    the key is the first component of the filename and the value is
    a list of the filenames with that component.'''
    keyfunc = lambda f: f.split('.', 1)[0]
    return dict( (k, list(g)) for k,g in groupby(
               sorted(filenames, key=keyfunc), key=keyfunc
           ) )

比如，给定这个列表

>>> test_data = [
...   exia.frame, exia.head, exia.swords, exia.legs,
...   exia.arms, exia.pilot, exia.gn_drive, lockon_stratos.data,
...   tieria_erde.data, ribbons_almark.data, otherstuff.dada
... ]

这个函数会产生

>>> group_by_name(test_data)
{'exia': ['exia.arms', 'exia.frame', 'exia.gn_drive', 'exia.head',
          'exia.legs', 'exia.pilot', 'exia.swords'],
 'lockon_stratos': ['lockon_stratos.data'],
 'otherstuff': ['otherstuff.dada'],
 'ribbons_almark': ['ribbons_almark.data'],
 'tieria_erde': ['tieria_erde.data']}

如果你想按后缀来索引文件名的话，只需要稍微修改一下就可以实现：

def group_by_extension(filenames):
    '''Puts the filenames in the given iterable into a dictionary where
    the key is the last component of the filename and the value is
    a list of the filenames with that extension.'''
    keyfunc = lambda f: f.split('.', 1)[1]
    return dict( (k, list(g)) for k,g in groupby(
               sorted(filenames, key=keyfunc), key=keyfunc
           ) )

唯一的不同在于 keyfunc = ... 这一行，我把关键字从 0 改成了 1。举个例子：

>>> group_by_extension(test_data)
{'arms': ['exia.arms'],
 'dada': ['otherstuff.dada'],
 'data': ['lockon_stratos.data', 'ribbons_almark.data', 'tieria_erde.data'],
 'frame': ['exia.frame'],
 'gn_drive': ['exia.gn_drive'],
 'head': ['exia.head'],
 'legs': ['exia.legs'],
 'pilot': ['exia.pilot'],
 'swords': ['exia.swords']}

不过，如果你想同时得到这两种分组，我觉得最好不要使用列表推导式，因为它只能处理一种方式，不能同时构建两个不同的字典。

from collections import defaultdict
def group_by_both(filenames):
    '''Puts the filenames in the given iterable into two dictionaries,
    where in the first, the key is the first component of the filename,
    and in the second, the key is the last component of the filename.
    The values in each dictionary are lists of the filenames with that
    base or extension.'''
    by_name = defaultdict(list)
    by_ext = defaultdict(list)
    for f in filenames:
        name, ext = f.split('.', 1)
        by_name[name] += [f]
        by_ext[ext] += [f]
    return by_name, by_ext