OILS / pea / pea_main.py View on Github | oilshell.org

424 lines, 234 significant
1#!/usr/bin/env python3
2"""
3pea_main.py
4
5A potential rewrite of mycpp.
6"""
7import ast
8from ast import AST, stmt, Module, ClassDef, FunctionDef, Assign
9import collections
10from dataclasses import dataclass
11import io
12import optparse
13import os
14import pickle
15from pprint import pprint
16import sys
17import time
18
19import typing
20from typing import Optional, Any
21
22from mycpp import pass_state
23
24
25START_TIME = time.time()
26
27def log(msg: str, *args: Any) -> None:
28 if args:
29 msg = msg % args
30 print('%.2f %s' % (time.time() - START_TIME, msg), file=sys.stderr)
31
32
33@dataclass
34class PyFile:
35 filename: str
36 namespace: str # C++ namespace
37 module: ast.Module # parsed representation
38
39
40class Program:
41 """A program is a collection of PyFiles."""
42
43 def __init__(self) -> None:
44 self.py_files : list[PyFile] = []
45
46 # As we parse, we add modules, and fill in the dictionaries with parsed
47 # types. Then other passes can retrieve the types with the same
48 # dictionaries.
49
50 # right now types are modules? Could change that
51 self.func_types: dict[FunctionDef, AST] = {}
52 self.method_types : dict[FunctionDef, AST] = {}
53 self.class_types : dict[ClassDef, Module] = {}
54 self.assign_types : dict[Assign, Module] = {}
55
56 # like mycpp: type and variable string. TODO: We shouldn't flatten it to a
57 # C type until later.
58 #
59 # Note: ImplPass parses the types. So I guess this could be limited to
60 # that?
61 # DoFunctionMethod() could make two passes?
62 # 1. collect vars
63 # 2. print code
64
65 self.local_vars : dict[FunctionDef, list[tuple[str, str]]] = {}
66
67 # ForwardDeclPass:
68 # OnMethod()
69 # OnSubclass()
70
71 # Then
72 # Calculate()
73 #
74 # PrototypesPass: # IsVirtual
75 self.virtual = pass_state.Virtual()
76
77 self.stats: dict[str, int] = {
78 # parsing stats
79 'num_files': 0,
80 'num_funcs': 0,
81 'num_classes': 0,
82 'num_methods': 0,
83 'num_assign': 0,
84
85 # ConstPass stats
86 'num_strings': 0,
87 }
88
89 def PrintStats(self) -> None:
90 pprint(self.stats, stream=sys.stderr)
91 print('', file=sys.stderr)
92
93
94class TypeSyntaxError(Exception):
95
96 def __init__(self, lineno: int, code_str: str):
97 self.lineno = lineno
98 self.code_str = code_str
99
100
101def ParseFiles(files: list[str], prog: Program) -> bool:
102
103 for filename in files:
104 with open(filename) as f:
105 contents = f.read()
106
107 try:
108 # Python 3.8+ supports type_comments=True
109 module = ast.parse(contents, filename=filename, type_comments=True)
110 except SyntaxError as e:
111 # This raises an exception for some reason
112 #e.print_file_and_line()
113 print('Error parsing %s: %s' % (filename, e))
114 return False
115
116 tmp = os.path.basename(filename)
117 namespace, _ = os.path.splitext(tmp)
118
119 prog.py_files.append(PyFile(filename, namespace, module))
120
121 prog.stats['num_files'] += 1
122
123 return True
124
125
126class ConstVisitor(ast.NodeVisitor):
127
128 def __init__(self, const_lookup: dict[str, int]):
129 ast.NodeVisitor.__init__(self)
130 self.const_lookup = const_lookup
131 self.str_id = 0
132
133 def visit_Constant(self, o: ast.Constant) -> None:
134 if isinstance(o.value, str):
135 self.const_lookup[o.value] = self.str_id
136 self.str_id += 1
137
138
139class ForwardDeclPass:
140 """Emit forward declarations."""
141 # TODO: Move this to ParsePass after comparing with mycpp.
142
143 def __init__(self, f: typing.IO[str]) -> None:
144 self.f = f
145
146 def DoPyFile(self, py_file: PyFile) -> None:
147
148 # TODO: could omit empty namespaces
149 namespace = py_file.namespace
150 self.f.write(f'namespace {namespace} {{ // forward declare\n')
151
152 for stmt in py_file.module.body:
153 match stmt:
154 case ClassDef():
155 class_name = stmt.name
156 self.f.write(f' class {class_name};\n')
157
158 self.f.write(f'}} // forward declare {namespace}\n')
159 self.f.write('\n')
160
161
162def _ParseFuncType(st: stmt) -> AST:
163 assert st.type_comment # caller checks this
164 try:
165 # This parses with the func_type production in the grammar
166 return ast.parse(st.type_comment, mode='func_type')
167 except SyntaxError:
168 raise TypeSyntaxError(st.lineno, st.type_comment)
169
170
171class PrototypesPass:
172 """Parse signatures and Emit function prototypes."""
173
174 def __init__(self, opts: Any, prog: Program, f: typing.IO[str]) -> None:
175 self.opts = opts
176 self.prog = prog
177 self.f = f
178
179 def DoClass(self, cls: ClassDef) -> None:
180 for stmt in cls.body:
181 match stmt:
182 case FunctionDef():
183 if stmt.type_comment:
184 sig = _ParseFuncType(stmt) # may raise
185
186 if self.opts.verbose:
187 print('METHOD')
188 print(ast.dump(sig, indent=' '))
189 # TODO: We need to print virtual here
190
191 self.prog.method_types[stmt] = sig # save for ImplPass
192 self.prog.stats['num_methods'] += 1
193
194 # TODO: assert that there aren't top-level statements?
195 case _:
196 pass
197
198 def DoPyFile(self, py_file: PyFile) -> None:
199 for stmt in py_file.module.body:
200 match stmt:
201 case FunctionDef():
202 if stmt.type_comment:
203 sig = _ParseFuncType(stmt) # may raise
204
205 if self.opts.verbose:
206 print('FUNC')
207 print(ast.dump(sig, indent=' '))
208
209 self.prog.func_types[stmt] = sig # save for ImplPass
210
211 self.prog.stats['num_funcs'] += 1
212
213 case ClassDef():
214 self.DoClass(stmt)
215 self.prog.stats['num_classes'] += 1
216
217 case _:
218 # Import, Assign, etc.
219 #print(stmt)
220
221 # TODO: omit __name__ == '__main__' etc.
222 # if __name__ == '__main__'
223 pass
224
225
226class ImplPass:
227 """Emit function and method bodies.
228
229 Algorithm:
230 collect local variables first
231 """
232
233 def __init__(self, prog: Program, f: typing.IO[str]) -> None:
234 self.prog = prog
235 self.f = f
236
237 # TODO: needs to be fully recursive, so you get bodies of loops, etc.
238 def DoBlock(self, stmts: list[stmt], indent: int=0) -> None:
239 """e.g. body of function, method, etc."""
240
241
242 #print('STMTS %s' % stmts)
243
244 ind_str = ' ' * indent
245
246 for stmt in stmts:
247 match stmt:
248 case Assign():
249 #print('%s* Assign' % ind_str)
250 #print(ast.dump(stmt, indent=' '))
251
252 if stmt.type_comment:
253 # This parses with the func_type production in the grammar
254 try:
255 typ = ast.parse(stmt.type_comment)
256 except SyntaxError as e:
257 # New syntax error
258 raise TypeSyntaxError(stmt.lineno, stmt.type_comment)
259
260 self.prog.assign_types[stmt] = typ
261
262 #print('%s TYPE: Assign' % ind_str)
263 #print(ast.dump(typ, indent=' '))
264
265 self.prog.stats['num_assign'] += 1
266
267 case _:
268 pass
269
270 def DoClass(self, cls: ClassDef) -> None:
271 for stmt in cls.body:
272 match stmt:
273 case FunctionDef():
274 self.DoBlock(stmt.body, indent=1)
275
276 case _:
277 pass
278
279 def DoPyFile(self, py_file: PyFile) -> None:
280 for stmt in py_file.module.body:
281 match stmt:
282 case ClassDef():
283 self.DoClass(stmt)
284
285 case FunctionDef():
286 self.DoBlock(stmt.body, indent=1)
287
288
289def Options() -> optparse.OptionParser:
290 """Returns an option parser instance."""
291
292 p = optparse.OptionParser()
293 p.add_option(
294 '-v', '--verbose', dest='verbose', action='store_true', default=False,
295 help='Show details about translation')
296
297 # Control which modules are exported to the header. Used by
298 # build/translate.sh.
299 p.add_option(
300 '--to-header', dest='to_header', action='append', default=[],
301 help='Export this module to a header, e.g. frontend.args')
302
303 p.add_option(
304 '--header-out', dest='header_out', default=None,
305 help='Write this header')
306
307 return p
308
309
310def main(argv: list[str]) -> int:
311
312 o = Options()
313 opts, argv = o.parse_args(argv)
314
315 action = argv[1]
316
317 # TODO: get rid of 'parse'
318 if action in ('parse', 'cpp'):
319 files = argv[2:]
320
321 # TODO:
322 # pass_state.Virtual
323 # this loops over functions and methods. But it has to be done BEFORE
324 # the PrototypesPass, or we need two passes. Gah!
325 # Could it be done in ConstVisitor? ConstVirtualVisitor?
326
327 # local_vars
328
329 prog = Program()
330 log('Pea begin')
331
332 if not ParseFiles(files, prog):
333 return 1
334 log('Parsed %d files and their type comments', len(files))
335 prog.PrintStats()
336
337 # This is the first pass
338
339 const_lookup: dict[str, int] = {}
340
341 v = ConstVisitor(const_lookup)
342 for py_file in prog.py_files:
343 v.visit(py_file.module)
344
345 log('Collected %d constants', len(const_lookup))
346
347 # TODO: respect header_out for these two passes
348 #out_f = sys.stdout
349 out_f = io.StringIO()
350
351 # ForwardDeclPass: module -> class
352 # TODO: Move trivial ForwardDeclPass into ParsePass, BEFORE constants,
353 # after comparing output with mycpp.
354 pass2 = ForwardDeclPass(out_f)
355 for py_file in prog.py_files:
356 namespace = py_file.namespace
357 pass2.DoPyFile(py_file)
358
359 log('Wrote forward declarations')
360 prog.PrintStats()
361
362 try:
363 # PrototypesPass: module -> class/method, func
364
365 pass3 = PrototypesPass(opts, prog, out_f)
366 for py_file in prog.py_files:
367 pass3.DoPyFile(py_file) # parses type comments in signatures
368
369 log('Wrote prototypes')
370 prog.PrintStats()
371
372 # ImplPass: module -> class/method, func; then probably a fully recursive thing
373
374 pass4 = ImplPass(prog, out_f)
375 for py_file in prog.py_files:
376 pass4.DoPyFile(py_file) # parses type comments in assignments
377
378 log('Wrote implementation')
379 prog.PrintStats()
380
381 except TypeSyntaxError as e:
382 log('Type comment syntax error on line %d of %s: %r',
383 e.lineno, py_file.filename, e.code_str)
384 return 1
385
386 log('Done')
387
388 elif action == 'dump-pickles':
389 files = argv[2:]
390
391 prog = Program()
392 log('Pea begin')
393
394 if not ParseFiles(files, prog):
395 return 1
396 log('Parsed %d files and their type comments', len(files))
397 prog.PrintStats()
398
399 # Note: can't use marshal here, because it only accepts simple types
400 pickle.dump(prog.py_files, sys.stdout.buffer)
401 log('Dumped pickle')
402
403 elif action == 'load-pickles':
404 while True:
405 try:
406 py_files = pickle.load(sys.stdin.buffer)
407 except EOFError:
408 break
409 log('Loaded pickle with %d files', len(py_files))
410
411 else:
412 raise RuntimeError('Invalid action %r' % action)
413
414 return 0
415
416
417if __name__ == '__main__':
418 try:
419 sys.exit(main(sys.argv))
420 except RuntimeError as e:
421 print('FATAL: %s' % e, file=sys.stderr)
422 sys.exit(1)
423
424# vim: sw=2