This instance will be upgraded to Heptapod 0.27.1 on 2021-12-08 at 22:00 UTC+1 (a few minutes of down time)

util.py 14.7 KB
Newer Older
1
2
3
4
"""Utilities for the analyses
=============================

"""
5
import re
6
from pathlib import Path
Pierre Augier's avatar
Pierre Augier committed
7
from textwrap import dedent
Pierre Augier's avatar
Pierre Augier committed
8

9
import gast as ast
Pierre Augier's avatar
Pierre Augier committed
10
11
import astunparse

12
from transonic.analyses import beniget
13
from transonic.analyses import extast
14
15


Pierre Augier's avatar
Pierre Augier committed
16
def print_dumped(source):
17
    """Pretty print the AST"""
Pierre Augier's avatar
Pierre Augier committed
18
    if isinstance(source, str):
19
        module = extast.parse(source)
Pierre Augier's avatar
Pierre Augier committed
20
21
22
23
24
25
        if len(module.body) == 1:
            node = module.body[0]
        else:
            node = module
    else:
        node = source
Pierre Augier's avatar
Pierre Augier committed
26
    print(astunparse.dump(node))
27
28


Pierre Augier's avatar
Pierre Augier committed
29
def print_unparsed(node):
30
    """Print the code corresponding to a tree or a node"""
31
    print(extast.unparse(node))
Pierre Augier's avatar
Pierre Augier committed
32
33


Pierre Augier's avatar
Pierre Augier committed
34
35
36
from transonic.analyses.capturex import CaptureX


37
38
39
40
41
42
43
44
45
46
def _fill_ast_annotations_function(function_def, ast_annotations):

    dict_node = ast_annotations.value
    for arg in function_def.args.args:
        if arg.annotation is not None:
            try:
                name = arg.arg
            except AttributeError:
                name = arg.id

Pierre Augier's avatar
Pierre Augier committed
47
            dict_node.keys.append(extast.Constant(value=name))
48
49
50
51
52
53
54
55
56
57
58
59
            dict_node.values.append(arg.annotation)


def _fill_ast_annotations_class(class_def, ast_annotations):

    dict_node = ast_annotations.value
    for node in class_def.body:
        if not isinstance(node, ast.AnnAssign):
            continue

        if node.annotation is not None:
            name = node.target.id
Pierre Augier's avatar
Pierre Augier committed
60
            dict_node.keys.append(extast.Constant(value=name))
61
62
63
64
            dict_node.values.append(node.annotation)


def get_annotations(object_def, namespace):
65
    """Create the annotations from a definition node"""
66
67
68
69

    # print_dump(object_def)

    ast_annotations = ast.Assign(
Pierre Augier's avatar
Pierre Augier committed
70
        targets=[extast.Name("annotations", ast.Store())],
71
72
73
74
75
76
77
78
79
80
81
82
        value=ast.Dict(keys=[], values=[]),
    )

    if isinstance(object_def, ast.FunctionDef):
        _fill_ast_annotations_function(object_def, ast_annotations)
    elif isinstance(object_def, ast.ClassDef):
        _fill_ast_annotations_class(object_def, ast_annotations)
    else:
        raise NotImplementedError

    # print_dump(ast_annotations)

83
    source = extast.unparse(ast_annotations)
84
85
86
87
88
89
90

    try:
        del namespace["__builtins__"]
    except KeyError:
        pass
    exec(source, namespace)
    return namespace["annotations"]
Pierre Augier's avatar
Pierre Augier committed
91
92
93


def filter_code_typevars(module, duc, ancestors):
94
    """Create a filtered code with what is needed to create the annotations"""
Pierre Augier's avatar
Pierre Augier committed
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

    module_filtered = ast.Module()
    kept = module_filtered.body = []
    suppressed = set()

    def fill_suppressed(def_):
        for user in def_.users():
            parent_in_body = ancestors.parents(user.node)[1]
            suppressed.add(parent_in_body)
            fill_suppressed(user)

    for node in module.body:
        if node in suppressed:
            continue

        if isinstance(node, ast.Import):
            if node.names[0].name in ["transonic", "numpy"]:
                kept.append(node)
            else:
                def_ = duc.chains[node.names[0]]
                fill_suppressed(def_)
            #     suppressed.add()
        elif isinstance(node, ast.ImportFrom):
            if node.module in ["transonic", "numpy"]:
                kept.append(node)

        elif isinstance(node, (ast.Assign, ast.AugAssign)):
            kept.append(node)

124
    return extast.unparse(module_filtered)
Pierre Augier's avatar
Pierre Augier committed
125
126
127


class AnalyseLines(ast.NodeVisitor):
128
129
    """Analyse to find the last line of a node"""

Pierre Augier's avatar
Pierre Augier committed
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
    def __init__(self, main_node):
        if isinstance(main_node, ast.Module):
            self.line_start = 1
        else:
            self.line_start = main_node.lineno
        self.line_last = self.line_start
        self.visit(main_node)

    def generic_visit(self, node):
        if hasattr(node, "lineno"):
            if node.lineno > self.line_last:
                self.line_last = node.lineno
        super().generic_visit(node)

    def get_lines(self):
        return self.line_start, self.line_last

    def get_code(self, source):
        lines = source.split("\n")

        stop = self.line_last
        nb_lines = len(lines)
        if nb_lines != stop:
            # find next not empty and not comment line
            ind = stop - 1
            next_line = ""
            while ind + 1 < nb_lines and next_line == "":
                ind += 1
                next_line = lines[ind].strip()
                if next_line.startswith("#"):
                    next_line = ""

            if any(
                next_line.startswith(character) for character in (")", "]", "}")
            ):
                stop = ind + 1

        return "\n".join(lines[self.line_start - 1 : stop])


def gather_rawcode_comments(node, code_module):
171
    """Get the comments in a node"""
Pierre Augier's avatar
Pierre Augier committed
172
173
    analysis = AnalyseLines(node)
    rawcode = dedent(analysis.get_code(code_module))
174
175
176
177
178
    comments = dedent(
        "\n".join(
            line for line in rawcode.split("\n") if line.strip().startswith("#")
        )
    )
Pierre Augier's avatar
Pierre Augier committed
179
    return rawcode, comments
180
181


182
packages_supported_by_pythran = [
183
184
185
    "bisect",
    "builtins",
    "cmath",
186
    "functools",
187
188
    "io",
    "itertools",
189
    "math",
190
191
192
193
    "operator",
    "random",
    "string",
    "time",
194
195
    "numpy",
    "numpy.fft",
196
197
    "numpy.linalg",
    "numpy.random",
198
    "scipy",
199
    "scipy.special",
200
]
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
201

Pierre Augier's avatar
Pierre Augier committed
202

203
204
205
def find_path(node: object, pathfile: str):
    """ Return the path of node (instance of ast.Import or ast.ImportFrom)
    """
Pierre Augier's avatar
Pierre Augier committed
206
    # FIXME find path in non local imports
207
208
209
210
211
    name = str()
    path = str()

    if isinstance(node, ast.ImportFrom):
        name = node.module
212
        if name in packages_supported_by_pythran:
213
214
            return None, None
        else:
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
215
216
            parent = Path(pathfile).parent
            path = parent / (str(name.replace(".", "/")) + ".py")
217
    else:
218
        if node.names[0].name in packages_supported_by_pythran:
219
220
221
222
223
224
225
            pass
        else:
            # TODO deal with an ast.Import
            raise NotImplementedError
    return name, path


226
def change_import_name(
227
    code_dep: str, changed_node: object, func_name: str, relative: str = None
228
):
Pierre Augier's avatar
Pierre Augier committed
229
230
    """Change the name of changed_node in code_dep by adding "__" + func + "__"
    at the beginning of the imported module, and return the modified code
231
232
233
234
235
    """
    mod = extast.parse(code_dep)
    for node in mod.body:
        if extast.unparse(node) == extast.unparse(changed_node):
            if isinstance(node, ast.ImportFrom):
Pierre Augier's avatar
Pierre Augier committed
236
                node.module = f"__ext__{func_name}__{node.module}"
237
            elif isinstance(node, ast.Import):
Pierre Augier's avatar
Pierre Augier committed
238
                node.names[0].name = f"__ext__{func_name}__{node.names[0].name}"
239
        if not relative:
240
            node.level = 0
241
    return extast.unparse(mod)
242
243
244


def filter_external_code(module: object, names: list):
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
245
    """ Filter the module to keep only the necessary nodes
246
247
248
        needed by functions or class in the parameter names
    """
    code_dependance_annotations = ""
249
    lines_code = []
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
    for node in module.body:
        for name in names:
            if isinstance(node, ast.FunctionDef):
                if node.name == extast.unparse(name).rstrip("\n\r"):
                    ancestors = beniget.Ancestors()
                    ancestors.visit(module)
                    duc = beniget.DefUseChains()
                    duc.visit(module)
                    udc = beniget.UseDefChains(duc)
                    capturex = CaptureX(
                        [node],
                        module,
                        ancestors,
                        defuse_chains=duc,
                        usedef_chains=udc,
                        consider_annotations=None,
                    )
267
                    lines_code.append(str(extast.unparse(node)))
268
269
270
                    code_dependance_annotations = capturex.make_code_external()
            if isinstance(node, ast.Assign):
                if node.targets[0].id == extast.unparse(name).rstrip("\n\r"):
271
                    lines_code.append(str(extast.unparse(node)))
272
273
            if isinstance(node, ast.ClassDef):
                if node.name == extast.unparse(name).rstrip("\n\r"):
274
                    lines_code.append(str(extast.unparse(node)))
275

276
    return code_dependance_annotations + "\n".join(lines_code)
277
278


279
def adapt_code_dependance(func: str, codes_dependance: str, jitted_dicts: dict):
280
281
282
283
    """
    Adapt code_dependance to the call of a jitted function in a jitted function:
        - Remove the import transonic
        - Remove the jitted function statement (i.e func = jit(func))
284
        - Add a import statement to the jitted function
285
286
        - remove the definition of the jitted function if its on the file, or remove the import statement
    """
287
    special = []
288
289
290
291
292
293
294
295
296
297
    module = extast.parse(codes_dependance)
    module_body = module.body.copy()
    jitted_functions = []
    for node in module_body:
        # remove the transonic import
        if isinstance(node, ast.ImportFrom):
            if node.module == "transonic":
                module.body.remove(node)
        # remove the jitted function by jit() (i.e. func = jit(func))
        # and add the import statement for the jitted function
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
        elif (
            isinstance(node, ast.Assign)
            and isinstance(node.value, ast.Call)
            and node.value.func.id == "jit"
        ):
            # if assigned name different from jitted function name
            if node.targets[0].id != node.value.args[0].id:
                # change function names in jitted dict
                # The list "special" contains functions that has to be write from jitted_dicts
                # see transonic.justintime:287
                def_func = extast.unparse(jitted_dicts["functions"][func])
                spl = re.split(r"(\W+)", def_func)
                spl = [
                    node.value.args[0].id if x == node.targets[0].id else x
                    for x in spl
                ]
                st = "".join(str(e) for e in spl)
                jitted_dicts["functions"][func] = extast.parse(st)
                special.append(func)
                jitted_functions.append(node.value.args[0].id)
            else:
                jitted_functions.append(node.targets[0].id)
            module.body.remove(node)
            module.body.insert(
                0,
                [
                    extast.parse(
                        "from "
                        + node.value.args[0].id
                        + " import "
                        + node.value.args[0].id
329
                    )
330
331
                ],
            )
332
333
334
335
336
337
338
339
340
341
342
343
344
345
    # remove the definition:
    for node in module_body:
        # of the jitted function in the file
        if isinstance(node, ast.FunctionDef):
            if node.name in jitted_functions:
                module.body.remove(node)
        # of the jitted imported function
        if isinstance(node, ast.ImportFrom):
            for name in node.names:
                if name.name in jitted_functions:
                    node.names.remove(name)
            # remove the importFrom if no function is imported
            if not node.names:
                module.body.remove(node)
346
    return extast.unparse(module), jitted_dicts, special, jitted_functions
347
348


Pierre Augier's avatar
Pierre Augier committed
349
code_ext = {"function": {}, "class": {}}
350
351
352
353
354
355
356
357


def get_exterior_code(
    codes_dependance: dict,
    pathfile: str,
    previous_file_name=None,
    classes: str = None,
    relative: bool = None,
358
    jitted_dicts: dict = None,
359
360
361
362
):
    """ Get all imported functions needed by boosted functions and methods at multiple levels,
        (i.e get functions needed by functions imported by boosted function) and add them into code_ext
    """
363
364
    special = []
    treated = []
365
    for func, dep in codes_dependance.items():
366
367
368
369
370
371
372
373
374
375
        if not dep:
            continue
        module_ext = extast.parse(dep)
        for node in module_ext.body:
            if not isinstance(node, (ast.ImportFrom, ast.Import)):
                continue
            # get the path of the imported module
            file_name, file_path = find_path(node, pathfile)
            # a jitted function or method needs another jitted function
            if file_name == "transonic":
Pierre Augier's avatar
Pierre Augier committed
376
377
378
379
380
381
                (
                    codes_dependance[func],
                    jitted_dicts,
                    spe,
                    treat,
                ) = adapt_code_dependance(
382
383
384
385
386
387
                    func, codes_dependance[func], jitted_dicts
                )
                # the "special" list signals that jitted functions has to be written
                # from jitted_dicts (see transonic/justintime.py:287)
                special = special + spe
                treated = treated + treat
388
389

    for func, dep in codes_dependance.items():
390
391
392
393
394
395
396
397
398
        if not dep:
            continue
        module_ext = extast.parse(dep)
        for node in module_ext.body:
            if not isinstance(node, (ast.ImportFrom, ast.Import)):
                continue
            # get the path of the imported module
            file_name, file_path = find_path(node, pathfile)
            # a jitted function or method needs another jitted function
Pierre Augier's avatar
Pierre Augier committed
399
            if not (file_name and file_name not in treated):
400
                continue
Pierre Augier's avatar
Pierre Augier committed
401
            new_file_name = f"__ext__{func}__{file_name}"
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
            # get the content of the file
            try:
                with open(str(file_path), "r") as file:
                    content = file.read()
            except:
                raise NotImplementedError(file_name + " can no be found")
            mod = extast.parse(content)
            # filter the code and add it to code_ext dict
            code_ext[classes][new_file_name] = str(
                filter_external_code(mod, node.names)
            )
            # change imported module names
            codes_dependance[func] = change_import_name(
                codes_dependance[func], node, func, relative
            )
            # recursively get the exterior codes
            if code_ext[classes][new_file_name]:
                get_exterior_code(
                    {func: code_ext[classes][new_file_name]},
                    pathfile,
                    new_file_name,
                    classes,
                )
                if previous_file_name:
                    code_ext[classes][previous_file_name] = change_import_name(
                        code_ext[classes][previous_file_name],
                        node,
                        func,
                        relative,
                    )
432
    return codes_dependance, code_ext, jitted_dicts, special
Pierre Augier's avatar
Pierre Augier committed
433
434
435
436
437
438
439
440
441
442
443
444
445


def extract_variable_annotations(fdef, namespace):
    variable_types = {}

    for node in fdef.body:
        if isinstance(node, ast.AnnAssign):
            name = node.target.id
            source = extast.unparse(node.annotation)
            result = eval(source, namespace)
            variable_types[name] = result

    return variable_types
Pierre Augier's avatar
Pierre Augier committed
446
447
448
449
450
451
452


def extract_returns_annotation(returns, namespace):
    if returns is None:
        return
    source = extast.unparse(returns)
    return eval(source, namespace)