7
7
import re
8
8
import time
9
9
import types
10
- import typing
10
+ from typing import Dict , FrozenSet , Tuple , TextIO
11
11
12
12
import umarshal
13
13
@@ -42,13 +42,14 @@ def get_localsplus(code: types.CodeType):
42
42
43
43
44
44
def get_localsplus_counts (code : types .CodeType ,
45
- names : tuple [str , ...],
46
- kinds : bytes ) -> tuple [int , int , int , int ]:
45
+ names : Tuple [str , ...],
46
+ kinds : bytes ) -> Tuple [int , int , int , int ]:
47
47
nlocals = 0
48
48
nplaincellvars = 0
49
49
ncellvars = 0
50
50
nfreevars = 0
51
- for name , kind in zip (names , kinds , strict = True ):
51
+ assert len (names ) == len (kinds )
52
+ for name , kind in zip (names , kinds ):
52
53
if kind & CO_FAST_LOCAL :
53
54
nlocals += 1
54
55
if kind & CO_FAST_CELL :
@@ -71,7 +72,7 @@ def get_localsplus_counts(code: types.CodeType,
71
72
PyUnicode_4BYTE_KIND = 4
72
73
73
74
74
- def analyze_character_width (s : str ) -> tuple [int , bool ]:
75
+ def analyze_character_width (s : str ) -> Tuple [int , bool ]:
75
76
maxchar = ' '
76
77
for c in s :
77
78
maxchar = max (maxchar , c )
@@ -86,12 +87,17 @@ def analyze_character_width(s: str) -> tuple[int, bool]:
86
87
return kind , ascii
87
88
88
89
90
+ def removesuffix (base : str , suffix : str ) -> str :
91
+ if base .endswith (suffix ):
92
+ return base [:len (base ) - len (suffix )]
93
+ return base
94
+
89
95
class Printer :
90
96
91
- def __init__ (self , file : typing . TextIO ):
97
+ def __init__ (self , file : TextIO ):
92
98
self .level = 0
93
99
self .file = file
94
- self .cache : dict [ tuple [type , object ], str ] = {}
100
+ self .cache : Dict [ Tuple [type , object ], str ] = {}
95
101
self .hits , self .misses = 0 , 0
96
102
self .patchups : list [str ] = []
97
103
self .write ('#include "Python.h"' )
@@ -231,7 +237,7 @@ def generate_code(self, name: str, code: types.CodeType) -> str:
231
237
# otherwise MSVC doesn't like it.
232
238
self .write (f".co_consts = { co_consts } ," )
233
239
self .write (f".co_names = { co_names } ," )
234
- self .write (f".co_firstinstr = (_Py_CODEUNIT *) { co_code . removesuffix ('.ob_base.ob_base' )} .ob_sval," )
240
+ self .write (f".co_firstinstr = (_Py_CODEUNIT *) { removesuffix (co_code , '.ob_base.ob_base' )} .ob_sval," )
235
241
self .write (f".co_exceptiontable = { co_exceptiontable } ," )
236
242
self .field (code , "co_flags" )
237
243
self .write (".co_warmup = QUICKENING_INITIAL_WARMUP_VALUE," )
@@ -259,7 +265,7 @@ def generate_code(self, name: str, code: types.CodeType) -> str:
259
265
self .write (f".co_freevars = { co_freevars } ," )
260
266
return f"& { name } .ob_base"
261
267
262
- def generate_tuple (self , name : str , t : tuple [object , ...]) -> str :
268
+ def generate_tuple (self , name : str , t : Tuple [object , ...]) -> str :
263
269
items = [self .generate (f"{ name } _{ i } " , it ) for i , it in enumerate (t )]
264
270
self .write ("static" )
265
271
with self .indent ():
@@ -323,7 +329,7 @@ def generate_complex(self, name: str, z: complex) -> str:
323
329
self .write (f".cval = {{ { z .real } , { z .imag } }}," )
324
330
return f"&{ name } .ob_base"
325
331
326
- def generate_frozenset (self , name : str , fs : frozenset [object ]) -> str :
332
+ def generate_frozenset (self , name : str , fs : FrozenSet [object ]) -> str :
327
333
ret = self .generate_tuple (name , tuple (sorted (fs )))
328
334
self .write ("// TODO: The above tuple should be a frozenset" )
329
335
return ret
@@ -336,34 +342,33 @@ def generate(self, name: str, obj: object) -> str:
336
342
# print(f"Cache hit {key!r:.40}: {self.cache[key]!r:.40}")
337
343
return self .cache [key ]
338
344
self .misses += 1
339
- match obj :
340
- case types .CodeType () | umarshal .Code () as code :
341
- val = self .generate_code (name , code )
342
- case tuple (t ):
343
- val = self .generate_tuple (name , t )
344
- case str (s ):
345
- val = self .generate_unicode (name , s )
346
- case bytes (b ):
347
- val = self .generate_bytes (name , b )
348
- case True :
349
- return "Py_True"
350
- case False :
351
- return "Py_False"
352
- case int (i ):
353
- val = self .generate_int (name , i )
354
- case float (x ):
355
- val = self .generate_float (name , x )
356
- case complex () as z :
357
- val = self .generate_complex (name , z )
358
- case frozenset (fs ):
359
- val = self .generate_frozenset (name , fs )
360
- case builtins .Ellipsis :
361
- return "Py_Ellipsis"
362
- case None :
363
- return "Py_None"
364
- case _:
365
- raise TypeError (
366
- f"Cannot generate code for { type (obj ).__name__ } object" )
345
+ if isinstance (obj , types .CodeType ) or isinstance (obj , umarshal .Code ):
346
+ val = self .generate_code (name , obj )
347
+ elif isinstance (obj , tuple ):
348
+ val = self .generate_tuple (name , obj )
349
+ elif isinstance (obj , str ):
350
+ val = self .generate_unicode (name , obj )
351
+ elif isinstance (obj , bytes ):
352
+ val = self .generate_bytes (name , obj )
353
+ elif obj is True :
354
+ return "Py_True"
355
+ elif obj is False :
356
+ return "Py_False"
357
+ elif isinstance (obj , int ):
358
+ val = self .generate_int (name , obj )
359
+ elif isinstance (obj , float ):
360
+ val = self .generate_float (name , obj )
361
+ elif isinstance (obj , complex ):
362
+ val = self .generate_complex (name , obj )
363
+ elif isinstance (obj , frozenset ):
364
+ val = self .generate_frozenset (name , obj )
365
+ elif obj is builtins .Ellipsis :
366
+ return "Py_Ellipsis"
367
+ elif obj is None :
368
+ return "Py_None"
369
+ else :
370
+ raise TypeError (
371
+ f"Cannot generate code for { type (obj ).__name__ } object" )
367
372
# print(f"Cache store {key!r:.40}: {val!r:.40}")
368
373
self .cache [key ] = val
369
374
return val
@@ -393,12 +398,12 @@ def decode_frozen_data(source: str) -> types.CodeType:
393
398
del lines [0 ]
394
399
while lines and re .match (FROZEN_DATA_LINE , lines [- 1 ]) is None :
395
400
del lines [- 1 ]
396
- values : tuple [int , ...] = ast .literal_eval ("" .join (lines ))
401
+ values : Tuple [int , ...] = ast .literal_eval ("" .join (lines ). strip ( ))
397
402
data = bytes (values )
398
403
return umarshal .loads (data )
399
404
400
405
401
- def generate (source : str , filename : str , modname : str , file : typing . TextIO ) -> None :
406
+ def generate (source : str , filename : str , modname : str , file : TextIO ) -> None :
402
407
if is_frozen_header (source ):
403
408
code = decode_frozen_data (source )
404
409
else :
@@ -439,7 +444,7 @@ def main() -> None:
439
444
verbose = args .verbose
440
445
with open (args .file , encoding = "utf-8" ) as f :
441
446
source = f .read ()
442
- modname = args .module or os .path .basename (args .file ). removesuffix ( ".py" )
447
+ modname = args .module or removesuffix ( os .path .basename (args .file ), ".py" )
443
448
output = args .output or modname + ".c"
444
449
with open (output , "w" , encoding = "utf-8" ) as file :
445
450
with report_time ("generate" ):
0 commit comments