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

"""
5
import re
6
import inspect
7
import beniget
8
from pathlib import Path
Pierre Augier's avatar
Pierre Augier committed
9
from textwrap import dedent
10
from typing import Callable
11
import gast as ast
12
from transonic.analyses import extast
13
from transonic.analyses.capturex import CaptureX
Pierre Augier's avatar
Pierre Augier committed
14
import astunparse
15
16


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


Pierre Augier's avatar
Pierre Augier committed
30
def print_unparsed(node):
31
    """Print the code corresponding to a tree or a node"""
32
    print(extast.unparse(node))
Pierre Augier's avatar
Pierre Augier committed
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
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):
63
    """Create the annotations from a definition node"""
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

    # 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)

81
    source = extast.unparse(ast_annotations)
82
83
84
85
86
87
88

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


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

    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)

122
    return extast.unparse(module_filtered)
Pierre Augier's avatar
Pierre Augier committed
123
124
125


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

Pierre Augier's avatar
Pierre Augier committed
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
    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):
169
    """Get the comments in a node"""
Pierre Augier's avatar
Pierre Augier committed
170
171
    analysis = AnalyseLines(node)
    rawcode = dedent(analysis.get_code(code_module))
172
173
174
175
176
    comments = dedent(
        "\n".join(
            line for line in rawcode.split("\n") if line.strip().startswith("#")
        )
    )
Pierre Augier's avatar
Pierre Augier committed
177
    return rawcode, comments
178
179


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

183
184
185
186
187
188
189
190
191
# FIXME find path in non local imports
def find_path(node: object, pathfile: str):
    """ Return the path of node (instance of ast.Import or ast.ImportFrom)
    """
    name = str()
    path = str()

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


206
def change_import_name(
207
    code_dep: str, changed_node: object, func_name: str, relative: str = None
208
):
Pierre Augier's avatar
Cleanup    
Pierre Augier committed
209
    """ Change the name of changed_node in code_dep by adding "__" + func + "__"
210
211
212
213
214
215
216
217
218
        at the beginning of the imported module, and return the modified code
    """
    mod = extast.parse(code_dep)
    for node in mod.body:
        if extast.unparse(node) == extast.unparse(changed_node):
            if isinstance(node, ast.ImportFrom):
                node.module = "__" + func_name + "__" + node.module
            elif isinstance(node, ast.Import):
                node.names[0].name = "__" + func_name + "__" + node.names[0].name
219
        if not relative:
220
            node.level = 0
221
    return extast.unparse(mod)
222
223
224


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

256
    return code_dependance_annotations + "\n".join(lines_code)
257
258


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


code_ext = {"function": {}, "classe": {}}


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

    for func, dep in codes_dependance.items():
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
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
        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 not (file_name and not file_name in treated):
                continue
            new_file_name = "__" + func + "__" + file_name
            # 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,
                    )
409
    return codes_dependance, code_ext, jitted_dicts, special
410
411
412
413
414
415
416
417
418
419
420
421
422
423


def get_source_with_numba(func: Callable):
    """Get the source and adapt to numba"""
    src = inspect.getsource(func)
    src = dedent(src)
    mod = extast.parse(src)
    for node in mod.body:
        node.decorator_list[0].keywords = [
            ast.keyword("nopython", value=ast.Str("True")),
            ast.keyword("nogil", value=ast.Str("True")),
        ]
    src = "\nfrom numba import jit \n" + extast.unparse(mod) + "\n"
    return src