-
-
Notifications
You must be signed in to change notification settings - Fork 2.7k
/
str_ops.py
105 lines (89 loc) · 3.3 KB
/
str_ops.py
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
"""Primitive str ops."""
from typing import List, Callable
from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER, EmitterInterface, EmitCallback
from mypyc.ir.rtypes import (
RType, object_rprimitive, str_rprimitive, bool_rprimitive, int_rprimitive, list_rprimitive
)
from mypyc.primitives.registry import (
func_op, binary_op, simple_emit, name_ref_op, method_op, call_emit, name_emit,
c_method_op, c_binary_op
)
# Get the 'str' type object.
name_ref_op('builtins.str',
result_type=object_rprimitive,
error_kind=ERR_NEVER,
emit=name_emit('&PyUnicode_Type', target_type='PyObject *'),
is_borrowed=True)
# str(obj)
func_op(name='builtins.str',
arg_types=[object_rprimitive],
result_type=str_rprimitive,
error_kind=ERR_MAGIC,
emit=call_emit('PyObject_Str'))
# str1 + str2
binary_op(op='+',
arg_types=[str_rprimitive, str_rprimitive],
result_type=str_rprimitive,
error_kind=ERR_MAGIC,
emit=call_emit('PyUnicode_Concat'))
# str.join(obj)
c_method_op(
name='join',
arg_types=[str_rprimitive, object_rprimitive],
return_type=str_rprimitive,
c_function_name='PyUnicode_Join',
error_kind=ERR_MAGIC
)
# str[index] (for an int index)
method_op(
name='__getitem__',
arg_types=[str_rprimitive, int_rprimitive],
result_type=str_rprimitive,
error_kind=ERR_MAGIC,
emit=call_emit('CPyStr_GetItem'))
# str.split(...)
str_split_types = [str_rprimitive, str_rprimitive, int_rprimitive] # type: List[RType]
str_split_emits = [simple_emit('{dest} = PyUnicode_Split({args[0]}, NULL, -1);'),
simple_emit('{dest} = PyUnicode_Split({args[0]}, {args[1]}, -1);'),
call_emit('CPyStr_Split')] \
# type: List[EmitCallback]
for i in range(len(str_split_types)):
method_op(
name='split',
arg_types=str_split_types[0:i+1],
result_type=list_rprimitive,
error_kind=ERR_MAGIC,
emit=str_split_emits[i])
# str1 += str2
#
# PyUnicodeAppend makes an effort to reuse the LHS when the refcount
# is 1. This is super dodgy but oh well, the interpreter does it.
c_binary_op(name='+=',
arg_types=[str_rprimitive, str_rprimitive],
return_type=str_rprimitive,
c_function_name='CPyStr_Append',
error_kind=ERR_MAGIC,
steals=[True, False])
def emit_str_compare(comparison: str) -> Callable[[EmitterInterface, List[str], str], None]:
def emit(emitter: EmitterInterface, args: List[str], dest: str) -> None:
temp = emitter.temp_name()
emitter.emit_declaration('int %s;' % temp)
emitter.emit_lines(
'%s = PyUnicode_Compare(%s, %s);' % (temp, args[0], args[1]),
'if (%s == -1 && PyErr_Occurred())' % temp,
' %s = 2;' % dest,
'else',
' %s = (%s %s);' % (dest, temp, comparison))
return emit
# str1 == str2
binary_op(op='==',
arg_types=[str_rprimitive, str_rprimitive],
result_type=bool_rprimitive,
error_kind=ERR_MAGIC,
emit=emit_str_compare('== 0'))
# str1 != str2
binary_op(op='!=',
arg_types=[str_rprimitive, str_rprimitive],
result_type=bool_rprimitive,
error_kind=ERR_MAGIC,
emit=emit_str_compare('!= 0'))