Correctly remove future_directives with # cython: language_level=2

2 files changed, 18 insertions, 12 deletions

diff --git a/Cython/Compiler/Main.py b/Cython/Compiler/Main.py
index 9130fbf19..48379d266 100644
--- a/Cython/Compiler/Main.py
+++ b/Cython/Compiler/Main.py
@@ -99,18 +99,17 @@ class Context(object):
 
     def set_language_level(self, level):
         from .Future import print_function, unicode_literals, absolute_import, division, generator_stop
-        future_directives = []
+        future_directives = set()
         if level == '3str':
-            future_directives = [print_function, absolute_import, division, generator_stop]
-            self.future_directives.discard(unicode_literals)
             level = 3
         else:
             level = int(level)
             if level >= 3:
-                future_directives = [print_function, unicode_literals, absolute_import, division, generator_stop]
+                future_directives.add(unicode_literals)
+        if level >= 3:
+            future_directives.update([print_function, absolute_import, division, generator_stop])
         self.language_level = level
-        if future_directives:
-            self.future_directives.update(future_directives)
+        self.future_directives = future_directives
         if level >= 3:
             self.modules['builtins'] = self.modules['__builtin__']
 
diff --git a/tests/run/language_level.srctree b/tests/run/language_level.srctree
index 91624982b..5e9ab640c 100644
--- a/tests/run/language_level.srctree
+++ b/tests/run/language_level.srctree
@@ -7,12 +7,19 @@ PYTHON -c "import infile2; import infile3"
 from Cython.Build.Dependencies import cythonize
 from distutils.core import setup
 
-setup(
-    ext_modules = (cythonize("infile*.py") +
-                   cythonize("directive2.py", compiler_directives={'language_level': 2}) +
-                   cythonize("directive3.py", compiler_directives={'language_level': 3})
-                   )
-)
+ext_modules = []
+
+# Test language_level specified in the cythonize() call
+ext_modules += cythonize("directive2.py", compiler_directives={'language_level': 2})
+ext_modules += cythonize("directive3.py", compiler_directives={'language_level': 3})
+
+# Test language_level specified in the source file. We give a
+# conflicting directive to cythonize() to check that the language_level
+# is correctly overridden when compiling
+ext_modules += cythonize("infile2.py", compiler_directives={'language_level': 3})
+ext_modules += cythonize("infile3.py", compiler_directives={'language_level': 2})
+
+setup(ext_modules=ext_modules)
 
 ######## directive3.py ########