Parse multi-line __future__ imports better (#519)

Augments #426
Augments #427
Augments #510

2 files changed, 49 insertions, 2 deletions

diff --git a/babel/util.py b/babel/util.py
index 996f902..af8c762 100644
--- a/babel/util.py
+++ b/babel/util.py
@@ -109,10 +109,23 @@ def parse_future_flags(fp, encoding='latin-1'):
     flags = 0
     try:
         body = fp.read().decode(encoding)
+
+        # Fix up the source to be (hopefully) parsable by regexpen.
+        # This will likely do untoward things if the source code itself is broken.
+
+        # (1) Fix `import (\n...` to be `import (...`.
+        body = re.sub(r'import\s*\([\r\n]+', 'import (', body)
+        # (2) Join line-ending commas with the next line.
+        body = re.sub(r',\s*[\r\n]+', ', ', body)
+        # (3) Remove backslash line continuations.
+        body = re.sub(r'\\\s*[\r\n]+', ' ', body)
+
         for m in PYTHON_FUTURE_IMPORT_re.finditer(body):
-            names = [x.strip() for x in m.group(1).split(',')]
+            names = [x.strip().strip('()') for x in m.group(1).split(',')]
             for name in names:
-                flags |= getattr(__future__, name).compiler_flag
+                feature = getattr(__future__, name, None)
+                if feature:
+                    flags |= feature.compiler_flag
     finally:
         fp.seek(pos)
     return flags
diff --git a/tests/test_util.py b/tests/test_util.py
index 068a102..4c769be 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -13,8 +13,11 @@
 
 import unittest
 
+import pytest
+
 from babel import util
 from babel._compat import BytesIO
+from babel.util import parse_future_flags
 
 
 def test_distinct():
@@ -69,3 +72,34 @@ def test_parse_encoding_undefined():
 
 def test_parse_encoding_non_ascii():
     assert parse_encoding(u'K\xf6ln') is None
+
+
+@pytest.mark.parametrize('source, result', [
+    ('''
+from __future__ import print_function,
+    division, with_statement,
+    unicode_literals
+''', 0x10000 | 0x2000 | 0x8000 | 0x20000),
+    ('''
+from __future__ import print_function, division
+print('hello')
+''', 0x10000 | 0x2000),
+    ('''
+from __future__ import print_function, division, unknown,,,,,
+print 'hello'
+''', 0x10000 | 0x2000),
+    ('''
+from __future__ import (
+    print_function,
+    division)
+''', 0x10000 | 0x2000),
+    ('''
+from __future__ import \\
+    print_function, \\
+    division
+''', 0x10000 | 0x2000),
+])
+def test_parse_future(source, result):
+    fp = BytesIO(source.encode('latin-1'))
+    flags = parse_future_flags(fp)
+    assert flags == result