util.py 14.6 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


Pierre Augier's avatar
Cleanup    
Pierre Augier committed
182
# TODO complete the list
183
184
185
186
187
188
packages_supported_by_pythran = [
    "functools",
    "math",
    "cmath",
    "numpy",
    "numpy.fft",
189
190
    "numpy.linalg",
    "numpy.random",
191
    "scipy",
192
    "scipy.special",
193
]
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
194

Pierre Augier's avatar
Pierre Augier committed
195

196
197
198
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
199
    # FIXME find path in non local imports
200
201
202
203
204
    name = str()
    path = str()

    if isinstance(node, ast.ImportFrom):
        name = node.module
205
        if name in packages_supported_by_pythran:
206
207
            return None, None
        else:
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
208
209
            parent = Path(pathfile).parent
            path = parent / (str(name.replace(".", "/")) + ".py")
210
    else:
211
        if node.names[0].name in packages_supported_by_pythran:
212
213
214
215
216
217
218
            pass
        else:
            # TODO deal with an ast.Import
            raise NotImplementedError
    return name, path


219
def change_import_name(
220
    code_dep: str, changed_node: object, func_name: str, relative: str = None
221
):
Pierre Augier's avatar
Pierre Augier committed
222
223
    """Change the name of changed_node in code_dep by adding "__" + func + "__"
    at the beginning of the imported module, and return the modified code
224
225
226
227
228
    """
    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
229
                node.module = f"__ext__{func_name}__{node.module}"
230
            elif isinstance(node, ast.Import):
Pierre Augier's avatar
Pierre Augier committed
231
                node.names[0].name = f"__ext__{func_name}__{node.names[0].name}"
232
        if not relative:
233
            node.level = 0
234
    return extast.unparse(mod)
235
236
237


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

269
    return code_dependance_annotations + "\n".join(lines_code)
270
271


272
def adapt_code_dependance(func: str, codes_dependance: str, jitted_dicts: dict):
273
274
275
276
    """
    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))
277
        - Add a import statement to the jitted function
278
279
        - remove the definition of the jitted function if its on the file, or remove the import statement
    """
280
    special = []
281
282
283
284
285
286
287
288
289
290
    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
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
319
320
321
        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
322
                    )
323
324
                ],
            )
325
326
327
328
329
330
331
332
333
334
335
336
337
338
    # 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)
339
    return extast.unparse(module), jitted_dicts, special, jitted_functions
340
341


Pierre Augier's avatar
Pierre Augier committed
342
code_ext = {"function": {}, "class": {}}
343
344
345
346
347
348
349
350


def get_exterior_code(
    codes_dependance: dict,
    pathfile: str,
    previous_file_name=None,
    classes: str = None,
    relative: bool = None,
351
    jitted_dicts: dict = None,
352
353
354
355
):
    """ 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
    """
356
357
    special = []
    treated = []
358
    for func, dep in codes_dependance.items():
359
360
361
362
363
364
365
366
367
368
        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
369
370
371
372
373
374
                (
                    codes_dependance[func],
                    jitted_dicts,
                    spe,
                    treat,
                ) = adapt_code_dependance(
375
376
377
378
379
380
                    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
381
382

    for func, dep in codes_dependance.items():
383
384
385
386
387
388
389
390
391
        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
392
            if not (file_name and file_name not in treated):
393
                continue
Pierre Augier's avatar
Pierre Augier committed
394
            new_file_name = f"__ext__{func}__{file_name}"
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
            # 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,
                    )
425
    return codes_dependance, code_ext, jitted_dicts, special
Pierre Augier's avatar
Pierre Augier committed
426
427
428
429
430
431
432
433
434
435
436
437
438


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
439
440
441
442
443
444
445


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