util.py 14.3 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 beniget
10
import gast as ast
Pierre Augier's avatar
Pierre Augier committed
11
12
import astunparse

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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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

            dict_node.keys.append(ast.Str(s=name))
            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
            dict_node.keys.append(ast.Str(s=name))
            dict_node.values.append(node.annotation)


def get_annotations(object_def, namespace):
65
    """Create the annotations from a definition node"""
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

    # print_dump(object_def)

    ast_annotations = ast.Assign(
        targets=[ast.Name(id="annotations", ctx=ast.Store(), annotation=None)],
        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


Pierre Augier's avatar
Cleanup    
Pierre Augier committed
182
# TODO complete the list
183
packages_supported_by_pythran = ["numpy", "math", "functools", "cmath", "scipy"]
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
184

Pierre Augier's avatar
Pierre Augier committed
185

186
187
188
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
189
    # FIXME find path in non local imports
190
191
192
193
194
    name = str()
    path = str()

    if isinstance(node, ast.ImportFrom):
        name = node.module
195
        if name in packages_supported_by_pythran:
196
197
            return None, None
        else:
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
198
199
            parent = Path(pathfile).parent
            path = parent / (str(name.replace(".", "/")) + ".py")
200
    else:
201
        if node.names[0].name in packages_supported_by_pythran:
202
203
204
205
206
207
208
            pass
        else:
            # TODO deal with an ast.Import
            raise NotImplementedError
    return name, path


209
def change_import_name(
210
    code_dep: str, changed_node: object, func_name: str, relative: str = None
211
):
Pierre Augier's avatar
Pierre Augier committed
212
213
    """Change the name of changed_node in code_dep by adding "__" + func + "__"
    at the beginning of the imported module, and return the modified code
214
215
216
217
218
    """
    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
219
                node.module = f"__ext__{func_name}__{node.module}"
220
            elif isinstance(node, ast.Import):
Pierre Augier's avatar
Pierre Augier committed
221
                node.names[0].name = f"__ext__{func_name}__{node.names[0].name}"
222
        if not relative:
223
            node.level = 0
224
    return extast.unparse(mod)
225
226
227


def filter_external_code(module: object, names: list):
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
228
    """ Filter the module to keep only the necessary nodes
229
230
231
        needed by functions or class in the parameter names
    """
    code_dependance_annotations = ""
232
    lines_code = []
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
    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,
                    )
250
                    lines_code.append(str(extast.unparse(node)))
251
252
253
                    code_dependance_annotations = capturex.make_code_external()
            if isinstance(node, ast.Assign):
                if node.targets[0].id == extast.unparse(name).rstrip("\n\r"):
254
                    lines_code.append(str(extast.unparse(node)))
255
256
            if isinstance(node, ast.ClassDef):
                if node.name == extast.unparse(name).rstrip("\n\r"):
257
                    lines_code.append(str(extast.unparse(node)))
258

259
    return code_dependance_annotations + "\n".join(lines_code)
260
261


262
def adapt_code_dependance(func: str, codes_dependance: str, jitted_dicts: dict):
263
264
265
266
    """
    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))
267
        - Add a import statement to the jitted function
268
269
        - remove the definition of the jitted function if its on the file, or remove the import statement
    """
270
    special = []
271
272
273
274
275
276
277
278
279
280
    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
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
        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
312
                    )
313
314
                ],
            )
315
316
317
318
319
320
321
322
323
324
325
326
327
328
    # 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)
329
    return extast.unparse(module), jitted_dicts, special, jitted_functions
330
331


Pierre Augier's avatar
Pierre Augier committed
332
code_ext = {"function": {}, "class": {}}
333
334
335
336
337
338
339
340


def get_exterior_code(
    codes_dependance: dict,
    pathfile: str,
    previous_file_name=None,
    classes: str = None,
    relative: bool = None,
341
    jitted_dicts: dict = None,
342
343
344
345
):
    """ 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
    """
346
347
    special = []
    treated = []
348
    for func, dep in codes_dependance.items():
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
        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":
                codes_dependance[
                    func
                ], jitted_dicts, spe, treat = adapt_code_dependance(
                    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
368
369

    for func, dep in codes_dependance.items():
370
371
372
373
374
375
376
377
378
        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
379
            if not (file_name and file_name not in treated):
380
                continue
Pierre Augier's avatar
Pierre Augier committed
381
            new_file_name = f"__ext__{func}__{file_name}"
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
            # 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,
                    )
412
    return codes_dependance, code_ext, jitted_dicts, special
Pierre Augier's avatar
Pierre Augier committed
413
414
415
416
417
418
419
420
421
422
423
424
425


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