1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
#!/usr/bin/env python
#
# This file is part of htcondor_utils.
#
# htcondor_utils is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# htcondor_utils is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with htcondor_utils. If not, see <http://www.gnu.org/licenses/>.
import argparse
import os
import fnmatch
import shutil
import tempfile
VERBOSE = False
def path_search(p, ext, strip_components):
filenames = []
p = os.path.abspath(p)
index = 1
for root, _, files in os.walk(p):
for f in files:
path = os.path.join(root, f)
if not os.path.isfile(path):
continue
if fnmatch.fnmatch(path, ext):
if strip_components:
head, tail = os.path.split(path)
strip_max = len(head.split(os.path.sep))
if strip_components <= strip_max:
components = head.split(os.path.sep)[strip_components:]
path = os.path.join(os.path.sep.join(components), tail)
else:
# strip_components = strip_max
print("Warning: Cannot strip {} components from {} (max {})".format(strip_components, path, strip_max))
if VERBOSE:
print("[{}]:{}".format(index, path))
filenames.append(path)
index += 1
return filenames
def generate_manifest(manifest):
dest = tempfile.NamedTemporaryFile("w+", delete=False)
for line in manifest:
dest.write(line + os.linesep)
return dest.name
def assign_chunks(manifest, chunks, output_dir, as_jobs=False):
label = 0
count = 0
written = 0
# max_chunks = 0
dest = None
output_dir = os.path.abspath(output_dir)
if VERBOSE:
print("Output directory: {}".format(output_dir))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
else:
shutil.rmtree(output_dir)
os.makedirs(output_dir)
if as_jobs:
manifest_count = len(file(manifest, 'r').readlines())
max_chunks = manifest_count / chunks
chunks = max_chunks
for line in open(manifest, 'r'):
if count % chunks == 0:
if dest:
dest.close()
if VERBOSE:
print("[{}]: {}, {} entries".format(label, os.path.basename(dest.name), written))
written = 0
dest = file(os.path.join(output_dir, "stdin." + str(label)), 'w')
label += 1
dest.write(line)
count += 1
written += 1
os.unlink(manifest)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--strip-components", default=0, type=int, help="Top-level directories to strip")
parser.add_argument("--output-dir", default="input")
parser.add_argument("--extension", default="*.*", type=str, help="Extension of files")
parser.add_argument("--chunks", default=1, type=int, help="Number of entries per file")
parser.add_argument("--as-jobs", action="store_true", help="Number of expected jobs")
parser.add_argument("--verbose", action="store_true")
parser.add_argument("search_path", action="store", type=str, help="Directory to search under")
args = parser.parse_args()
global VERBOSE
VERBOSE = args.verbose
chunks = args.chunks
if chunks < 1:
chunks = 1
index_at = 0
files = path_search(args.search_path, args.extension, args.strip_components)
if files:
index_at = 1
if VERBOSE:
print("{} files".format(len(files)))
if not files:
exit(0)
manifest = generate_manifest(files)
assign_chunks(manifest, chunks, args.output_dir, as_jobs=args.as_jobs)
if __name__ == "__main__":
main()
|
