Repository: pyccel/pyccel
Pull request #231 (2020-01-24 12:01:16)

Import preprocess
Allow all 4 import syntaxes:
"from module import func1, func2"
"from module import func1 as f1, func2 as f2"
"import module"
"import module as mod"


Patch:
diff --git a/pyccel/ast/utilities.py b/pyccel/ast/utilities.py
index 6a43af62..9e10f896 100644
--- a/pyccel/ast/utilities.py
+++ b/pyccel/ast/utilities.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from sympy.core.function import Application
-from .core import DottedName
+from .core import AsName
 from .core import Import
 from .core import Range, Len , Enumerate, Zip, Product, Map
 from .core import FunctionDef, Return, Assign
@@ -141,7 +141,7 @@ def builtin_function(expr, args=None):
     return None
 
 # TODO add documentation
-builtin_import_registery = ('numpy', 'scipy', 'itertools', 'math')
+builtin_import_registery = ('numpy', 'scipy.constants', 'itertools', 'math')
 
 def builtin_import(expr):
     """Returns a builtin pyccel-extension function/object from an import."""
@@ -152,45 +152,40 @@ def builtin_import(expr):
     if expr.source is None:
         return []
 
-    source = expr.source
-    if isinstance(source, DottedName):
-        source = source.name[0]
-    else:
-        source = str(source)
+    source = str(expr.source)
 
         # TODO imrove
     imports = []
-    for i in range(len(expr.target)):
+    for target in expr.target:
+        if isinstance(target, AsName):
+            import_name = target.target
+            code_name = target.name
+        else:
+            import_name = str(target)
+            code_name = import_name
         if source == 'numpy':
 
-            target = str(expr.target[i])
-            if target in numpy_functions.keys():
-                imports.append((target, numpy_functions[target]))
+            if import_name in numpy_functions.keys():
+                imports.append((code_name, numpy_functions[import_name]))
 
-            elif target in math_functions.keys():
-                imports.append((target, math_functions[target]))
+            elif import_name in math_functions.keys():
+                imports.append((code_name, math_functions[import_name]))
 
-            elif target in numpy_constants.keys():
-                imports.append((target, numpy_constants[target]))
+            elif import_name in numpy_constants.keys():
+                imports.append((code_name, numpy_constants[import_name]))
 
         elif source == 'math':
 
-            target = str(expr.target[i])
-
-            if target in math_functions.keys():
-                imports.append((target, math_functions[target]))
+            if import_name in math_functions.keys():
+                imports.append((code_name, math_functions[import_name]))
 
-        elif source == 'scipy':
-            # TODO improve: source must be scipy.constants
-            #      - use dynamic import?
-            target = str(expr.target[i])
-            if target in scipy_constants.keys():
-                imports.append((target, scipy_constants[target]))
+        elif source == 'scipy.constants':
+            if import_name in scipy_constants.keys():
+                imports.append((code_name, scipy_constants[import_name]))
         elif source == 'itertools':
-            target = str(expr.target[i])
 
-            if target == 'product':
-                imports.append((target, Product))
+            if import_name == 'product':
+                imports.append((code_name, Product))
 
 
 
diff --git a/pyccel/parser/base.py b/pyccel/parser/base.py
index a53bb964..a53f75ba 100644
--- a/pyccel/parser/base.py
+++ b/pyccel/parser/base.py
@@ -54,7 +54,7 @@ from redbaron import YieldNode
 from redbaron import YieldAtomNode
 from redbaron import BreakNode, ContinueNode
 from redbaron import GetitemNode, SliceNode
-from redbaron import ImportNode, FromImportNode
+from redbaron import FromImportNode
 from redbaron import DottedAsNameNode, DecoratorNode
 from redbaron import NameAsNameNode
 from redbaron import LambdaNode
diff --git a/pyccel/parser/messages.py b/pyccel/parser/messages.py
index c7d91fb9..cee63b03 100644
--- a/pyccel/parser/messages.py
+++ b/pyccel/parser/messages.py
@@ -73,6 +73,7 @@ PYCCEL_MISSING_HEADER = 'Cannot find associated header'
 MACRO_MISSING_HEADER_OR_FUNC = 'Cannot find associated header/FunctionDef to macro'
 PYCCEL_UNFOUND_IMPORTED_MODULE = 'Unable to import'
 FOUND_DUPLICATED_IMPORT = 'Duplicated import '
+PYCCEL_UNEXPECTED_IMPORT = 'Pyccel has not correctly handled "import module" statement. Try again with "from module import function" syntax'
 
 IMPORTING_EXISTING_IDENTIFIED = \
         'Trying to import an identifier that already exists in the namespace. Hint: use import as'
diff --git a/pyccel/parser/syntactic.py b/pyccel/parser/syntactic.py
index 3804911c..daca4072 100644
--- a/pyccel/parser/syntactic.py
+++ b/pyccel/parser/syntactic.py
@@ -54,7 +54,7 @@ from redbaron import YieldNode
 from redbaron import YieldAtomNode
 from redbaron import BreakNode, ContinueNode
 from redbaron import GetitemNode, SliceNode
-from redbaron import ImportNode, FromImportNode
+from redbaron import FromImportNode
 from redbaron import DottedAsNameNode, DecoratorNode
 from redbaron import NameAsNameNode
 from redbaron import LambdaNode
@@ -126,7 +126,7 @@ from pyccel.ast.datatypes import sp_dtype, str_dtype
 
 
 from pyccel.parser.utilities import omp_statement, acc_statement
-from pyccel.parser.utilities import fst_move_directives
+from pyccel.parser.utilities import fst_move_directives, preprocess_imports
 from pyccel.parser.utilities import reconstruct_pragma_multilines
 from pyccel.parser.utilities import is_valid_filename_pyh, is_valid_filename_py
 from pyccel.parser.utilities import read_file
@@ -231,6 +231,7 @@ class SyntaxParser(BasicParser):
             errors.check()
             raise SystemExit(0)
 
+        preprocess_imports(red)
         red = fst_move_directives(red)
         self._fst = red
 
@@ -441,17 +442,10 @@ class SyntaxParser(BasicParser):
             return Symbol(val)
 
     def _visit_ImportNode(self, stmt):
-        if not isinstance(stmt.parent, (RedBaron, DefNode)):
-            errors.report(PYCCEL_RESTRICTION_IMPORT,
-                          bounding_box=stmt.absolute_bounding_box,
-                          severity='error')
-
-        if isinstance(stmt.parent, DefNode):
-            errors.report(PYCCEL_RESTRICTION_IMPORT_IN_DEF,
-                          bounding_box=stmt.absolute_bounding_box,
-                          severity='error')
+        errors.report(PYCCEL_UNEXPECTED_IMPORT,
+                      bounding_box=stmt.absolute_bounding_box,
+                      severity='error')
 
-        # in an import statement, we can have seperate target by commas
         ls = self._visit(stmt.value)
         ls = get_default_path(ls)
         expr = Import(ls)
diff --git a/pyccel/parser/utilities.py b/pyccel/parser/utilities.py
index 395acaab..f0ea6ab1 100644
--- a/pyccel/parser/utilities.py
+++ b/pyccel/parser/utilities.py
@@ -4,7 +4,8 @@
 
 from redbaron import (CommentNode, ForNode, DefNode, WithNode,
                       IfNode, ElseNode, ElifNode, IfelseblockNode,
-                      EndlNode)
+                      EndlNode, DottedAsNameNode, NameNode,
+                      CallNode, RedBaron, AtomtrailersNode)
 
 from sympy import srepr
 from pyccel.ast import DottedName
@@ -12,6 +13,9 @@ from sympy import Symbol
 from sympy.printing.dot import dotprint
 import os
 
+import string
+import random
+
 pyccel_external_lib = {"mpi4py"             : "pyccel.stdlib.external.mpi4py",
                        "scipy.linalg.lapack": "pyccel.stdlib.external.lapack",
                        "scipy.linalg.blas"  : "pyccel.stdlib.external.blas",
@@ -21,6 +25,14 @@ pyccel_external_lib = {"mpi4py"             : "pyccel.stdlib.external.mpi4py",
                        "numpy.linalg"       : "numpy",
                        "scipy.interpolate._fitpack":"pyccel.stdlib.external.fitpack"}
 
+#==============================================================================
+def random_string( n ):
+    # we remove uppercase letters because of f2py
+    chars    = string.ascii_lowercase + string.digits
+    selector = random.SystemRandom()
+    return ''.join( selector.choice( chars ) for _ in range( n ) )
+
+#==============================================================================
 
 def read_file(filename):
     """Returns the source code from a filename."""
@@ -191,6 +203,134 @@ def fst_move_directives(x):
     return x
 # ...
 
+def get_module_name( dotted_as_node ):
+    code_name = dotted_as_node.target
+    if (code_name != ""):
+        return [ code_name ]
+    else:
+        import_name = dotted_as_node.value
+        return import_name.dumps().split('.')
+
+def preprocess_imports(red):
+    imports = red.find_all("import", recursive = False)
+    code_names = []
+    idx = 0
+    while idx < len(imports):
+        imp_node = imports[idx]
+
+        if (len(imp_node)>1):
+            for imp in imp_node:
+                # Create and insert a line containing 1 import
+                node_to_add = RedBaron("import "+imp.dumps())[0]
+                red.insert(imp_node.index_on_parent, node_to_add)
+                imports.insert(idx, node_to_add)
+                code_names.append(get_module_name(node_to_add.value[0]))
+                idx+=1
+
+            # remove the multiple import line
+            red.remove(imp_node)
+            imports.remove(imp_node)
+        else:
+            code_names.append(get_module_name(imp_node[0]))
+            idx+=1
+
+    handled = set()
+
+    for idx,imp_node in enumerate(imports):
+        code_name_str = '.'.join(code_names[idx])
+        if (code_name_str not in handled):
+            handled.add(code_name_str)
+            line = find_import_usage(imp_node.value[0], red)
+            imp_node.replace(RedBaron(line)[0])
+
+def get_occurences_in_code(code_name, scope):
+    if (len(code_name)==1):
+        return scope.find_all(["dotted_as_name","name"],
+                              lambda occ: (isinstance(occ,DottedAsNameNode) and occ.target == code_name[0])
+                                                or (isinstance(occ,NameNode) and occ.value == code_name[0]))
+    else:
+        return scope.find_all(["name"], lambda occ: (isinstance(occ,NameNode) and occ.value == code_name[0]))
+
+def find_import_usage(stmt, scope, usage = None):
+    code_name = get_module_name(stmt)
+    import_name = stmt.value
+    if (len(code_name)==1):
+        if (code_name[0] == import_name.dumps()):
+            my_node = import_name[0]
+        else:
+            my_node = stmt
+    else:
+        my_node = import_name[0]
+
+    code_replacement = '_'.join(code_name) + "_" + random_string(6)
+
+    if (usage is None):
+        usage = get_occurences_in_code(code_name, scope)
+
+    i = 0
+
+    # ignore occurences before import
+    while (usage[i] != my_node):
+        i+=1
+    i+=1
+
+    if (i>=len(usage)):
+        return "\n"
+
+    targets = set()
+
+    while (i < len(usage)):
+        imp_node = usage[i].parent_find("import")
+
+        # If occurs in an import node
+        if (imp_node is not None and get_module_name(imp_node.value[0])==code_name):
+            imp_node = imp_node
+            occ_scope = usage[i].parent
+            while occ_scope.parent != None and not isinstance(occ_scope, (RedBaron, DefNode)):
+                occ_scope = occ_scope.parent
+
+            # if in the same scope then break
+            if (occ_scope == scope):
+                replacement_line = find_import_usage(imp_node.value[0], occ_scope, usage[i:])
+                imp_node.replace(RedBaron(replacement_line)[0])
+                break
+            else:
+                # if in a function then skip over the rest of the function
+                occ_usage = get_occurences_in_code(code_name, occ_scope)
+
+                i+=len(occ_usage)
+
+                replacement_line = find_import_usage(imp_node.value[0], occ_scope, occ_usage)
+                imp_node.replace(RedBaron(replacement_line)[0])
+        else:
+            phrase = usage[i].parent
+            pos_in_phrase = usage[i].index_on_parent
+
+            # Collect the object called from the import module
+            if isinstance(phrase, AtomtrailersNode) and pos_in_phrase==0 and (len(code_name) < len(phrase.value)):
+                match = (len(phrase)== len(code_name)+1 or (len(phrase) == len(code_name)+2 and isinstance(phrase[-1], CallNode)))
+                for j in range(1, len(code_name)):
+                    if (phrase.value[j].value!=code_name[j]):
+                        match = False
+                        break
+                if (match):
+                    func = phrase.value[pos_in_phrase+len(code_name)]
+                    if isinstance(func, NameNode):
+                        func_name = func.dumps()
+                        targets.add(func_name+" as "+code_replacement+"_"+func_name)
+
+                        line = phrase.dumps()
+                        new_line = line.replace('.'.join(code_name)+'.',code_replacement+'_')
+                        phrase.replace(RedBaron(new_line)[0])
+
+            i+=1
+
+    if (len(targets)!=0):
+        return "from "+import_name.dumps()+" import "+", ".join(targets)
+    else:
+        return "\n"
+
+
 # ...
 def reconstruct_pragma_multilines(header):
     """Must be called once we visit an annotated comment, to get the remaining


Test patch:
diff --git a/tests/pyccel/scripts/import_syntax/collisions.py b/tests/pyccel/scripts/import_syntax/collisions.py
new file mode 100644
index 00000000..cc06b258
--- /dev/null
+++ b/tests/pyccel/scripts/import_syntax/collisions.py
@@ -0,0 +1,24 @@
+# pylint: disable=reimported
+import numpy as np
+a = np.ones(3)
+
+def f():
+    import math as np
+    x = np.sqrt(3)
+    print(x)
+
+b = np.zeros(3)
+
+np = a[0]+b[0]
+
+print(np)
+
+f()
+
+import math as np
+x = np.sqrt(3)
+print(x)
+
+import scipy.constants
+
+print(scipy.constants.pi)
diff --git a/tests/pyccel/scripts/import_syntax/from_mod_import.py b/tests/pyccel/scripts/import_syntax/from_mod_import.py
new file mode 100644
index 00000000..b987da76
--- /dev/null
+++ b/tests/pyccel/scripts/import_syntax/from_mod_import.py
@@ -0,0 +1,44 @@
+from numpy import zeros, shape
+
+#$ header function matmat(double [:,:], double [:,:], double [:,:])
+def matmat(a,b,c):
+    nm = shape(a)
+    mp = shape(b)
+
+    n = nm[0]
+    m = nm[1]
+    p = mp[1]
+
+    for i in range(0, n):
+        for j in range(0, p):
+            for k in range(0, m):
+                c[i,j] = c[i,j] + a[i,k]*b[k,j]
+
+n = 3
+m = 4
+p = 3
+
+a = zeros((n,m), 'double')
+b = zeros((m,p), 'double')
+
+for i in range(0, n):
+    for j in range(0, m):
+        a[i,j] = (i-j)*1.0
+        print(a[i,j])
+    print()
+print()
+
+for i in range(0, m):
+    for j in range(0, p):
+        b[i,j] = (i+j)*1.0
+        print(b[i,j])
+    print()
+print()
+
+c = zeros((n,p),'double')
+matmat(a,b,c)
+
+for i in range(0, n):
+    for j in range(0, p):
+        print(c[i,j])
+    print()
diff --git a/tests/pyccel/scripts/import_syntax/from_mod_import_as.py b/tests/pyccel/scripts/import_syntax/from_mod_import_as.py
new file mode 100644
index 00000000..a1507b20
--- /dev/null
+++ b/tests/pyccel/scripts/import_syntax/from_mod_import_as.py
@@ -0,0 +1,45 @@
+from numpy import zeros as zilch
+from numpy import shape as form
+
+#$ header function matmat(double [:,:], double [:,:], double [:,:])
+def matmat(a,b,c):
+    nm = form(a)
+    mp = form(b)
+
+    n = nm[0]
+    m = nm[1]
+    p = mp[1]
+
+    for i in range(0, n):
+        for j in range(0, p):
+            for k in range(0, m):
+                c[i,j] = c[i,j] + a[i,k]*b[k,j]
+
+n = 3
+m = 4
+p = 3
+
+a = zilch((n,m), 'double')
+b = zilch((m,p), 'double')
+
+for i in range(0, n):
+    for j in range(0, m):
+        a[i,j] = (i-j)*1.0
+        print(a[i,j])
+    print()
+print()
+
+for i in range(0, m):
+    for j in range(0, p):
+        b[i,j] = (i+j)*1.0
+        print(b[i,j])
+    print()
+print()
+
+c = zilch((n,p),'double')
+matmat(a,b,c)
+
+for i in range(0, n):
+    for j in range(0, p):
+        print(c[i,j])
+    print()
diff --git a/tests/pyccel/scripts/import_syntax/import_mod.py b/tests/pyccel/scripts/import_syntax/import_mod.py
new file mode 100644
index 00000000..c7799dc6
--- /dev/null
+++ b/tests/pyccel/scripts/import_syntax/import_mod.py
@@ -0,0 +1,44 @@
+import numpy
+
+#$ header function matmat(double [:,:], double [:,:], double [:,:])
+def matmat(a,b,c):
+    nm = numpy.shape(a)
+    mp = numpy.shape(b)
+
+    n = nm[0]
+    m = nm[1]
+    p = mp[1]
+
+    for i in range(0, n):
+        for j in range(0, p):
+            for k in range(0, m):
+                c[i,j] = c[i,j] + a[i,k]*b[k,j]
+
+n = 3
+m = 4
+p = 3
+
+a = numpy.zeros((n,m), 'double')
+b = numpy.zeros((m,p), 'double')
+
+for i in range(0, n):
+    for j in range(0, m):
+        a[i,j] = (i-j)*1.0
+        print(a[i,j])
+    print()
+print()
+
+for i in range(0, m):
+    for j in range(0, p):
+        b[i,j] = (i+j)*1.0
+        print(b[i,j])
+    print()
+print()
+
+c = numpy.zeros((n,p),'double')
+matmat(a,b,c)
+
+for i in range(0, n):
+    for j in range(0, p):
+        print(c[i,j])
+    print()
diff --git a/tests/pyccel/scripts/import_syntax/import_mod_as.py b/tests/pyccel/scripts/import_syntax/import_mod_as.py
new file mode 100644
index 00000000..6140c462
--- /dev/null
+++ b/tests/pyccel/scripts/import_syntax/import_mod_as.py
@@ -0,0 +1,44 @@
+import numpy as np
+
+#$ header function matmat(double [:,:], double [:,:], double [:,:])
+def matmat(a,b,c):
+    nm = np.shape(a)
+    mp = np.shape(b)
+
+    n = nm[0]
+    m = nm[1]
+    p = mp[1]
+
+    for i in range(0, n):
+        for j in range(0, p):
+            for k in range(0, m):
+                c[i,j] = c[i,j] + a[i,k]*b[k,j]
+
+n = 3
+m = 4
+p = 3
+
+a = np.zeros((n,m), 'double')
+b = np.zeros((m,p), 'double')
+
+for i in range(0, n):
+    for j in range(0, m):
+        a[i,j] = (i-j)*1.0
+        print(a[i,j])
+    print()
+print()
+
+for i in range(0, m):
+    for j in range(0, p):
+        b[i,j] = (i+j)*1.0
+        print(b[i,j])
+    print()
+print()
+
+c = np.zeros((n,p),'double')
+matmat(a,b,c)
+
+for i in range(0, n):
+    for j in range(0, p):
+        print(c[i,j])
+    print()
diff --git a/tests/pyccel/test_pyccel.py b/tests/pyccel/test_pyccel.py
index aa072bba..411c7f52 100644
--- a/tests/pyccel/test_pyccel.py
+++ b/tests/pyccel/test_pyccel.py
@@ -244,3 +244,12 @@ def test_in_specified():
                                         ] )
 def test_hope_benchmarks( test_file ):
     pyccel_test(test_file)
+
+@pytest.mark.parametrize( "test_file", ["scripts/import_syntax/from_mod_import.py",
+                                        "scripts/import_syntax/from_mod_import_as.py",
+                                        "scripts/import_syntax/import_mod.py",
+                                        "scripts/import_syntax/import_mod_as.py",
+                                        "scripts/import_syntax/collisions.py"
+                                        ] )
+def test_import_syntax( test_file ):
+    pyccel_test(test_file)