1"""
2This script processes the output from the C preprocessor and extracts all
3qstr. Each qstr is transformed into a qstr definition of the form 'Q(...)'.
4
5This script works with Python 2.6, 2.7, 3.3 and 3.4.
6"""
7
8from __future__ import print_function
9
10import io
11import os
12import re
13import subprocess
14import sys
15import multiprocessing, multiprocessing.dummy
16
17
18# Extract MP_QSTR_FOO macros.
19_MODE_QSTR = "qstr"
20
21# Extract MP_COMPRESSED_ROM_TEXT("") macros.  (Which come from MP_ERROR_TEXT)
22_MODE_COMPRESS = "compress"
23
24
25def preprocess():
26    if any(src in args.dependencies for src in args.changed_sources):
27        sources = args.sources
28    elif any(args.changed_sources):
29        sources = args.changed_sources
30    else:
31        sources = args.sources
32    csources = []
33    cxxsources = []
34    for source in sources:
35        if source.endswith(".cpp"):
36            cxxsources.append(source)
37        elif source.endswith(".c"):
38            csources.append(source)
39    try:
40        os.makedirs(os.path.dirname(args.output[0]))
41    except OSError:
42        pass
43
44    def pp(flags):
45        def run(files):
46            return subprocess.check_output(args.pp + flags + files)
47
48        return run
49
50    try:
51        cpus = multiprocessing.cpu_count()
52    except NotImplementedError:
53        cpus = 1
54    p = multiprocessing.dummy.Pool(cpus)
55    with open(args.output[0], "wb") as out_file:
56        for flags, sources in (
57            (args.cflags, csources),
58            (args.cxxflags, cxxsources),
59        ):
60            batch_size = (len(sources) + cpus - 1) // cpus
61            chunks = [sources[i : i + batch_size] for i in range(0, len(sources), batch_size or 1)]
62            for output in p.imap(pp(flags), chunks):
63                out_file.write(output)
64
65
66def write_out(fname, output):
67    if output:
68        for m, r in [("/", "__"), ("\\", "__"), (":", "@"), ("..", "@@")]:
69            fname = fname.replace(m, r)
70        with open(args.output_dir + "/" + fname + "." + args.mode, "w") as f:
71            f.write("\n".join(output) + "\n")
72
73
74def process_file(f):
75    re_line = re.compile(r"#[line]*\s\d+\s\"([^\"]+)\"")
76    if args.mode == _MODE_QSTR:
77        re_match = re.compile(r"MP_QSTR_[_a-zA-Z0-9]+")
78    elif args.mode == _MODE_COMPRESS:
79        re_match = re.compile(r'MP_COMPRESSED_ROM_TEXT\("([^"]*)"\)')
80    output = []
81    last_fname = None
82    for line in f:
83        if line.isspace():
84            continue
85        # match gcc-like output (# n "file") and msvc-like output (#line n "file")
86        if line.startswith(("# ", "#line")):
87            m = re_line.match(line)
88            assert m is not None
89            fname = m.group(1)
90            if os.path.splitext(fname)[1] not in [".c", ".cpp"]:
91                continue
92            if fname != last_fname:
93                write_out(last_fname, output)
94                output = []
95                last_fname = fname
96            continue
97        for match in re_match.findall(line):
98            if args.mode == _MODE_QSTR:
99                name = match.replace("MP_QSTR_", "")
100                output.append("Q(" + name + ")")
101            elif args.mode == _MODE_COMPRESS:
102                output.append(match)
103
104    if last_fname:
105        write_out(last_fname, output)
106    return ""
107
108
109def cat_together():
110    import glob
111    import hashlib
112
113    hasher = hashlib.md5()
114    all_lines = []
115    outf = open(args.output_dir + "/out", "wb")
116    for fname in glob.glob(args.output_dir + "/*." + args.mode):
117        with open(fname, "rb") as f:
118            lines = f.readlines()
119            all_lines += lines
120    all_lines.sort()
121    all_lines = b"\n".join(all_lines)
122    outf.write(all_lines)
123    outf.close()
124    hasher.update(all_lines)
125    new_hash = hasher.hexdigest()
126    # print(new_hash)
127    old_hash = None
128    try:
129        with open(args.output_file + ".hash") as f:
130            old_hash = f.read()
131    except IOError:
132        pass
133    mode_full = "QSTR"
134    if args.mode == _MODE_COMPRESS:
135        mode_full = "Compressed data"
136    if old_hash != new_hash:
137        print(mode_full, "updated")
138        try:
139            # rename below might fail if file exists
140            os.remove(args.output_file)
141        except:
142            pass
143        os.rename(args.output_dir + "/out", args.output_file)
144        with open(args.output_file + ".hash", "w") as f:
145            f.write(new_hash)
146    else:
147        print(mode_full, "not updated")
148
149
150if __name__ == "__main__":
151    if len(sys.argv) < 6:
152        print("usage: %s command mode input_filename output_dir output_file" % sys.argv[0])
153        sys.exit(2)
154
155    class Args:
156        pass
157
158    args = Args()
159    args.command = sys.argv[1]
160
161    if args.command == "pp":
162        named_args = {
163            s: []
164            for s in [
165                "pp",
166                "output",
167                "cflags",
168                "cxxflags",
169                "sources",
170                "changed_sources",
171                "dependencies",
172            ]
173        }
174
175        for arg in sys.argv[1:]:
176            if arg in named_args:
177                current_tok = arg
178            else:
179                named_args[current_tok].append(arg)
180
181        if not named_args["pp"] or len(named_args["output"]) != 1:
182            print("usage: %s %s ..." % (sys.argv[0], " ... ".join(named_args)))
183            sys.exit(2)
184
185        for k, v in named_args.items():
186            setattr(args, k, v)
187
188        preprocess()
189        sys.exit(0)
190
191    args.mode = sys.argv[2]
192    args.input_filename = sys.argv[3]  # Unused for command=cat
193    args.output_dir = sys.argv[4]
194    args.output_file = None if len(sys.argv) == 5 else sys.argv[5]  # Unused for command=split
195
196    if args.mode not in (_MODE_QSTR, _MODE_COMPRESS):
197        print("error: mode %s unrecognised" % sys.argv[2])
198        sys.exit(2)
199
200    try:
201        os.makedirs(args.output_dir)
202    except OSError:
203        pass
204
205    if args.command == "split":
206        with io.open(args.input_filename, encoding="utf-8") as infile:
207            process_file(infile)
208
209    if args.command == "cat":
210        cat_together()
211