• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import glob
2import os
3import os.path
4
5# XXX need tests:
6# * walk_tree()
7# * glob_tree()
8# * iter_files_by_suffix()
9
10
11C_SOURCE_SUFFIXES = ('.c', '.h')
12
13
14def _walk_tree(root, *,
15               _walk=os.walk,
16               ):
17    # A wrapper around os.walk that resolves the filenames.
18    for parent, _, names in _walk(root):
19        for name in names:
20            yield os.path.join(parent, name)
21
22
23def walk_tree(root, *,
24              suffix=None,
25              walk=_walk_tree,
26              ):
27    """Yield each file in the tree under the given directory name.
28
29    If "suffix" is provided then only files with that suffix will
30    be included.
31    """
32    if suffix and not isinstance(suffix, str):
33        raise ValueError('suffix must be a string')
34
35    for filename in walk(root):
36        if suffix and not filename.endswith(suffix):
37            continue
38        yield filename
39
40
41def glob_tree(root, *,
42              suffix=None,
43              _glob=glob.iglob,
44              _escape=glob.escape,
45              _join=os.path.join,
46              ):
47    """Yield each file in the tree under the given directory name.
48
49    If "suffix" is provided then only files with that suffix will
50    be included.
51    """
52    suffix = suffix or ''
53    if not isinstance(suffix, str):
54        raise ValueError('suffix must be a string')
55
56    for filename in _glob(_join(_escape(root), f'*{suffix}')):
57        yield filename
58    for filename in _glob(_join(_escape(root), f'**/*{suffix}')):
59        yield filename
60
61
62def iter_files(root, suffix=None, relparent=None, *,
63               get_files=None,
64               _glob=glob_tree,
65               _walk=walk_tree,
66               ):
67    """Yield each file in the tree under the given directory name.
68
69    If "root" is a non-string iterable then do the same for each of
70    those trees.
71
72    If "suffix" is provided then only files with that suffix will
73    be included.
74
75    if "relparent" is provided then it is used to resolve each
76    filename as a relative path.
77    """
78    if get_files is None:
79        get_files = os.walk
80    if not isinstance(root, str):
81        roots = root
82        for root in roots:
83            yield from iter_files(root, suffix, relparent,
84                                  get_files=get_files,
85                                  _glob=_glob, _walk=_walk)
86        return
87
88    # Use the right "walk" function.
89    if get_files in (glob.glob, glob.iglob, glob_tree):
90        get_files = _glob
91    else:
92        _files = _walk_tree if get_files in (os.walk, walk_tree) else get_files
93        get_files = (lambda *a, **k: _walk(*a, walk=_files, **k))
94
95    # Handle a single suffix.
96    if suffix and not isinstance(suffix, str):
97        filenames = get_files(root)
98        suffix = tuple(suffix)
99    else:
100        filenames = get_files(root, suffix=suffix)
101        suffix = None
102
103    for filename in filenames:
104        if suffix and not isinstance(suffix, str):  # multiple suffixes
105            if not filename.endswith(suffix):
106                continue
107        if relparent:
108            filename = os.path.relpath(filename, relparent)
109        yield filename
110
111
112def iter_files_by_suffix(root, suffixes, relparent=None, *,
113                         walk=walk_tree,
114                         _iter_files=iter_files,
115                         ):
116    """Yield each file in the tree that has the given suffixes.
117
118    Unlike iter_files(), the results are in the original suffix order.
119    """
120    if isinstance(suffixes, str):
121        suffixes = [suffixes]
122    # XXX Ignore repeated suffixes?
123    for suffix in suffixes:
124        yield from _iter_files(root, suffix, relparent)
125