Merge pull request #95 from skeledrew/fix-string-exec

Fix source string exec

birdseye/utils.py

@@ -1,11 +1,14 @@
 import ast
+import io
 import json
+import linecache
 import ntpath
 import os
 import sys
 import token
 import types
 from sys import version_info
+
 from littleutils import strip_required_prefix
 
 # noinspection PyUnreachableCode
@@ -165,46 +168,28 @@ def default(self, o):
             return method()
 
 
-try:
-
-    # Python 3
-    from tokenize import open as open_with_encoding_check
-
-except ImportError:
-
-    # Python 2
-    from lib2to3.pgen2.tokenize import detect_encoding
-    import io
-
-
-    def open_with_encoding_check(filename):  # type: ignore
-        """Open a file in read only mode using the encoding detected by
-        detect_encoding().
-        """
-        fp = io.open(filename, 'rb')
-        try:
-            encoding, lines = detect_encoding(fp.readline)
-            fp.seek(0)
-            text = io.TextIOWrapper(fp, encoding, line_buffering=True)
-            text.mode = 'r'
-            return text
-        except:
-            fp.close()
-            raise
-
-
 def read_source_file(filename):
-    from lib2to3.pgen2.tokenize import cookie_re
-
-    if filename.endswith('.pyc'):
-        filename = filename[:-1]
-
-    with open_with_encoding_check(filename) as f:
-        return ''.join([
-            '\n' if i < 2 and cookie_re.match(line)
-            else line
-            for i, line in enumerate(f)
-        ])
+    if PY3:
+        from tokenize import detect_encoding, cookie_re
+    else:
+        from lib2to3.pgen2.tokenize import detect_encoding, cookie_re
+
+    lines = linecache.getlines(filename)
+    text = ''.join(lines)
+
+    if not isinstance(text, Text):
+        encoding = detect_encoding(io.BytesIO(text).readline)[0]
+        text = text.decode(encoding)  # noqa
+        lines = [line.decode(encoding) for line in lines]
+
+    # In python 2 it's a syntax error to parse unicode
+    # with an encoding declaration, so we remove it but
+    # leave empty lines in its place to keep line numbers the same
+    return ''.join([
+        '\n' if i < 2 and cookie_re.match(line)
+        else line
+        for i, line in enumerate(lines)
+    ])
 
 
 def source_without_decorators(tokens, function_node):

tests/test_utils.py

@@ -1,6 +1,7 @@
 # coding=utf8
 
 import ast
+import linecache
 import unittest
 from tempfile import mkstemp
 
@@ -99,6 +100,7 @@ def test_open_with_encoding_check(self):
         def write(stuff):
             with open(filename, 'wb') as f:
                 f.write(stuff)
+            linecache.cache.pop(filename, None)
 
         def read():
             return read_source_file(filename).strip()
@@ -146,6 +148,30 @@ def test_cheap_repr(self):
         series = df[0]
         self.assertEqual(cheap_repr(series), "0 = 0; 1 = 100; 2 = 200; ...; 97 = 9700; 98 = 9800; 99 = 9900")
 
+    def test_read_source_file_as_string(self):
+        import linecache
+        import birdseye
+
+        source = """
+from birdseye import eye
+
+@eye
+def func():
+    return 42
+"""
+        filename = "<some-non-filename-string>"
+        co = compile(source, filename, "exec")
+        linecache.cache[filename] = (
+            len(source),
+            None,
+            [line + '\n' for line in source.splitlines()],
+            filename,
+        )
+        exec(co, {"birdseye": birdseye})
+        rv = read_source_file(filename)
+        self.assertEqual(rv, source)
+        return
+
 
 if __name__ == '__main__':
     unittest.main()