zmc
2023-08-08 e792e9a60d958b93aef96050644f369feb25d61b
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
# -*- coding: utf-8 -*-
"""
Tools for searching bytecode for key statements that indicate the need for additional resources, such as data files
and package metadata.
 
By *bytecode* I mean the ``code`` object given by ``compile()``, accessible from the ``__code__`` attribute of any
non-builtin function or, in PyInstallerLand, the ``PyiModuleGraph.node("some.module").code`` attribute. The best
guide for bytecode format I have found is the disassembler reference: https://docs.python.org/3/library/dis.html
 
This parser implementation aims to combine the flexibility and speed of regex with the clarity of the output of
``dis.dis(code)``. It has not achieved the 2nd, but C'est la vie...
 
The biggest clarity killer here is the ``EXTENDED_ARG`` opcode which can appear almost anywhere and therefore needs
to be tiptoed around at every step. If this code needs to expand significantly, I would recommend an upgrade to a
regex-based grammar parsing library such as Reparse. This way, little steps like unpacking ``EXTENDED_ARGS`` can be
defined once then simply referenced forming a nice hierarchy rather than copied everywhere its needed.
"""
 
import dis
import re
from types import CodeType
from typing import Pattern
 
from PyInstaller import compat
 
# opcode name -> opcode map
# Python 3.11 introduced specialized opcodes that are not covered by opcode.opmap (and equivalent dis.opmap), but dis
# has a private map of all opcodes called _all_opmap. So use the latter, if available.
opmap = getattr(dis, '_all_opmap', dis.opmap)
 
 
def _instruction_to_regex(x: str):
    """
    Get a regex-escaped opcode byte from its human readable name.
    """
    return re.escape(bytes([opmap[x]]))
 
 
def bytecode_regex(pattern: bytes, flags=re.VERBOSE | re.DOTALL):
    """
    A regex-powered Python bytecode matcher.
 
    ``bytecode_regex`` provides a very thin wrapper around :func:`re.compile`.
 
      * Any opcode names wrapped in backticks are substituted for their corresponding opcode bytes.
      * Patterns are compiled in VERBOSE mode by default so that whitespace and comments may be used.
 
    This aims to mirror the output of :func:`dis.dis`, which is far more readable than looking at raw byte strings.
    """
    assert isinstance(pattern, bytes)
 
    # Replace anything wrapped in backticks with regex-escaped opcodes.
    pattern = re.sub(
        rb"`(\w+)`",
        lambda m: _instruction_to_regex(m[1].decode()),
        pattern,
    )
    return re.compile(pattern, flags=flags)
 
 
def finditer(pattern: Pattern, string: bytes):
    """
    Call ``pattern.finditer(string)``, but remove any matches beginning on an odd byte (i.e., matches where
    match.start() is not a multiple of 2).
 
    This should be used to avoid false positive matches where a bytecode pair's argument is mistaken for an opcode.
    """
    assert isinstance(string, bytes)
    string = _cleanup_bytecode_string(string)
    matches = pattern.finditer(string)
    while True:
        for match in matches:
            if match.start() % 2 == 0:
                # All is good. This match starts on an OPCODE.
                yield match
            else:
                # This match has started on an odd byte, meaning that it is a false positive and should be skipped.
                # There is a very slim chance that a genuine match overlaps this one and, because re.finditer() does not
                # allow overlapping matches, it would be lost. To avoid that, restart the regex scan, starting at the
                # next even byte.
                matches = pattern.finditer(string, match.start() + 1)
                break
        else:
            break
 
 
# Opcodes involved in function calls with constant arguments. The differences between python versions are handled by
# variables below, which are then used to construct the _call_function_bytecode regex.
# NOTE1: the _OPCODES_* entries are typically used in (non-capturing) groups that match the opcode plus an arbitrary
# argument. But because the entries themselves may contain more than on opcode (with OR operator between them), they
# themselves need to be enclosed in another (non-capturing) group. E.g., "(?:(?:_OPCODES_FUNCTION_GLOBAL).)".
# NOTE2: _OPCODES_EXTENDED_ARG2 is an exception, as it is used as a list of opcodes to exclude, i.e.,
# "[^_OPCODES_EXTENDED_ARG2]". Therefore, multiple opcodes are not separated by the OR operator.
if not compat.is_py37:
    _OPCODES_EXTENDED_ARG = rb"`EXTENDED_ARG`"
    _OPCODES_EXTENDED_ARG2 = _OPCODES_EXTENDED_ARG
    _OPCODES_FUNCTION_GLOBAL = rb"`LOAD_NAME`|`LOAD_GLOBAL`|`LOAD_FAST`"
    _OPCODES_FUNCTION_LOAD = rb"`LOAD_ATTR`"
    _OPCODES_FUNCTION_ARGS = rb"`LOAD_CONST`"
    _OPCODES_FUNCTION_CALL = rb"`CALL_FUNCTION`|`CALL_FUNCTION_EX`"
 
    def _cleanup_bytecode_string(bytecode):
        return bytecode  # Nothing to do here
elif not compat.is_py311:
    # Python 3.7 introduced two new function-related opcodes, LOAD_METHOD and CALL_METHOD
    _OPCODES_EXTENDED_ARG = rb"`EXTENDED_ARG`"
    _OPCODES_EXTENDED_ARG2 = _OPCODES_EXTENDED_ARG
    _OPCODES_FUNCTION_GLOBAL = rb"`LOAD_NAME`|`LOAD_GLOBAL`|`LOAD_FAST`"
    _OPCODES_FUNCTION_LOAD = rb"`LOAD_ATTR`|`LOAD_METHOD`"
    _OPCODES_FUNCTION_ARGS = rb"`LOAD_CONST`"
    _OPCODES_FUNCTION_CALL = rb"`CALL_FUNCTION`|`CALL_METHOD`|`CALL_FUNCTION_EX`"
 
    def _cleanup_bytecode_string(bytecode):
        return bytecode  # Nothing to do here
else:
    # Python 3.11 removed CALL_FUNCTION and CALL_METHOD, and replaced them with PRECALL + CALL instruction sequence.
    # As both PRECALL and CALL have the same parameter (the argument count), we need to match only up to the PRECALL.
    # The CALL_FUNCTION_EX is still present.
    # From Python 3.11b1 on, there is an EXTENDED_ARG_QUICK specialization opcode present.
    _OPCODES_EXTENDED_ARG = rb"`EXTENDED_ARG`|`EXTENDED_ARG_QUICK`"
    _OPCODES_EXTENDED_ARG2 = rb"`EXTENDED_ARG``EXTENDED_ARG_QUICK`"  # Special case; see note above the if/else block!
    _OPCODES_FUNCTION_GLOBAL = rb"`LOAD_NAME`|`LOAD_GLOBAL`|`LOAD_FAST`"
    _OPCODES_FUNCTION_LOAD = rb"`LOAD_ATTR`|`LOAD_METHOD`"
    _OPCODES_FUNCTION_ARGS = rb"`LOAD_CONST`"
    _OPCODES_FUNCTION_CALL = rb"`PRECALL`|`CALL_FUNCTION_EX`"
 
    # Starting with python 3.11, the bytecode is peppered with CACHE instructions (which dis module conveniently hides
    # unless show_caches=True is used). Dealing with these CACHE instructions in regex rules is going to render them
    # unreadable, so instead we pre-process the bytecode and filter the offending opcodes out.
    _cache_instruction_filter = bytecode_regex(rb"(`CACHE`.)|(..)")
 
    def _cleanup_bytecode_string(bytecode):
        return _cache_instruction_filter.sub(rb"\2", bytecode)
 
 
# language=PythonVerboseRegExp
_call_function_bytecode = bytecode_regex(
    rb"""
    # Matches `global_function('some', 'constant', 'arguments')`.
 
    # Load the global function. In code with >256 of names, this may require extended name references.
    (
     (?:(?:""" + _OPCODES_EXTENDED_ARG + rb""").)*
     (?:(?:""" + _OPCODES_FUNCTION_GLOBAL + rb""").)
    )
 
    # For foo.bar.whizz(), the above is the 'foo', below is the 'bar.whizz' (one opcode per name component, each
    # possibly preceded by name reference extension).
    (
     (?:
       (?:(?:""" + _OPCODES_EXTENDED_ARG + rb""").)*
       (?:""" + _OPCODES_FUNCTION_LOAD + rb""").
     )*
    )
 
    # Load however many arguments it takes. These (for now) must all be constants.
    # Again, code with >256 constants may need extended enumeration.
    (
      (?:
        (?:(?:""" + _OPCODES_EXTENDED_ARG + rb""").)*
        (?:""" + _OPCODES_FUNCTION_ARGS + rb""").
      )*
    )
 
    # Call the function. If opcode is CALL_FUNCTION_EX, the parameter are flags. For other opcodes, the parameter
    # is the argument count (which may be > 256).
    (
      (?:(?:""" + _OPCODES_EXTENDED_ARG + rb""").)*
      (?:""" + _OPCODES_FUNCTION_CALL + rb""").
    )
"""
)
 
# language=PythonVerboseRegExp
_extended_arg_bytecode = bytecode_regex(
    rb"""(
    # Arbitrary number of EXTENDED_ARG pairs.
    (?:(?:""" + _OPCODES_EXTENDED_ARG + rb""").)*
 
    # Followed by some other instruction (usually a LOAD).
    [^""" + _OPCODES_EXTENDED_ARG2 + rb"""].
)"""
)
 
 
def extended_arguments(extended_args: bytes):
    """
    Unpack the (extended) integer used to reference names or constants.
 
    The input should be a bytecode snippet of the following form::
 
        EXTENDED_ARG    ?      # Repeated 0-4 times.
        LOAD_xxx        ?      # Any of LOAD_NAME/LOAD_CONST/LOAD_METHOD/...
 
    Each ? byte combined together gives the number we want.
    """
    return int.from_bytes(extended_args[1::2], "big")
 
 
def load(raw: bytes, code: CodeType) -> str:
    """
    Parse an (extended) LOAD_xxx instruction.
    """
    # Get the enumeration.
    index = extended_arguments(raw)
 
    # Work out what that enumeration was for (constant/local var/global var).
 
    # If the last instruction byte is a LOAD_FAST:
    if raw[-2] == opmap["LOAD_FAST"]:
        # Then this is a local variable.
        return code.co_varnames[index]
    # Or if it is a LOAD_CONST:
    if raw[-2] == opmap["LOAD_CONST"]:
        # Then this is a literal.
        return code.co_consts[index]
    # Otherwise, it is a global name.
    if raw[-2] == opmap["LOAD_GLOBAL"] and compat.is_py311:
        # In python 3.11, namei>>1 is pushed on stack...
        return code.co_names[index >> 1]
    return code.co_names[index]
 
 
def loads(raw: bytes, code: CodeType) -> list:
    """
    Parse multiple consecutive LOAD_xxx instructions. Or load() in a for loop.
 
    May be used to unpack a function's parameters or nested attributes ``(foo.bar.pop.whack)``.
    """
    return [load(i, code) for i in _extended_arg_bytecode.findall(raw)]
 
 
def function_calls(code: CodeType) -> list:
    """
    Scan a code object for all function calls on constant arguments.
    """
    match: re.Match
    out = []
 
    for match in finditer(_call_function_bytecode, code.co_code):
        function_root, methods, args, function_call = match.groups()
 
        # For foo():
        #   `function_root` contains 'foo' and `methods` is empty.
        # For foo.bar.whizz():
        #   `function_root` contains 'foo' and `methods` contains the rest.
        function_root = load(function_root, code)
        methods = loads(methods, code)
        function = ".".join([function_root] + methods)
 
        args = loads(args, code)
        if function_call[0] == opmap['CALL_FUNCTION_EX']:
            flags = extended_arguments(function_call)
            if flags != 0:
                # Keyword arguments present. Unhandled at the moment.
                continue
            # In calls with const arguments, args contains a single
            # tuple with all values.
            if len(args) != 1 or not isinstance(args[0], tuple):
                continue
            args = list(args[0])
        else:
            arg_count = extended_arguments(function_call)
 
            if arg_count != len(args):
                # This happens if there are variable or keyword arguments. Bail out in either case.
                continue
 
        out.append((function, args))
 
    return out
 
 
def search_recursively(search: callable, code: CodeType, _memo=None) -> dict:
    """
    Apply a search function to a code object, recursing into child code objects (function definitions).
    """
    if _memo is None:
        _memo = {}
    if code not in _memo:
        _memo[code] = search(code)
        for const in code.co_consts:
            if isinstance(const, CodeType):
                search_recursively(search, const, _memo)
    return _memo
 
 
def recursive_function_calls(code: CodeType) -> dict:
    """
    Scan a code object for function calls on constant arguments, recursing into function definitions and bodies of
    comprehension loops.
    """
    return search_recursively(function_calls, code)
 
 
def any_alias(full_name: str):
    """List possible aliases of a fully qualified Python name.
 
        >>> list(any_alias("foo.bar.wizz"))
        ['foo.bar.wizz', 'bar.wizz', 'wizz']
 
    This crudely allows us to capture uses of wizz() under any of
    ::
        import foo
        foo.bar.wizz()
    ::
        from foo import bar
        bar.wizz()
    ::
        from foo.bar import wizz
        wizz()
 
    However, it will fail for any form of aliases and quite likely find false matches.
    """
    parts = full_name.split('.')
    while parts:
        yield ".".join(parts)
        parts = parts[1:]