evaluate.site: copy/adapt site-packages related functionality from stdlib

jedi/evaluate/site.py

@@ -0,0 +1,110 @@
+"""An adapted copy of relevant site-packages functionality from Python stdlib.
+
+This file contains some functions related to handling site-packages in Python
+with jedi-specific modifications:
+
+- the functions operate on sys_path argument rather than global sys.path
+
+- in .pth files "import ..." lines that allow execution of arbitrary code are
+  skipped to prevent code injection into jedi interpreter
+
+"""
+
+# Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010,
+# 2011, 2012, 2013, 2014, 2015 Python Software Foundation; All Rights Reserved
+
+from __future__ import print_function
+
+import sys
+import os
+
+
+def makepath(*paths):
+    dir = os.path.join(*paths)
+    try:
+        dir = os.path.abspath(dir)
+    except OSError:
+        pass
+    return dir, os.path.normcase(dir)
+
+
+def _init_pathinfo(sys_path):
+    """Return a set containing all existing directory entries from sys_path"""
+    d = set()
+    for dir in sys_path:
+        try:
+            if os.path.isdir(dir):
+                dir, dircase = makepath(dir)
+                d.add(dircase)
+        except TypeError:
+            continue
+    return d
+
+
+def addpackage(sys_path, sitedir, name, known_paths):
+    """Process a .pth file within the site-packages directory:
+       For each line in the file, either combine it with sitedir to a path
+       and add that to known_paths, or execute it if it starts with 'import '.
+    """
+    if known_paths is None:
+        known_paths = _init_pathinfo(sys_path)
+        reset = 1
+    else:
+        reset = 0
+    fullname = os.path.join(sitedir, name)
+    try:
+        f = open(fullname, "r")
+    except OSError:
+        return
+    with f:
+        for n, line in enumerate(f):
+            if line.startswith("#"):
+                continue
+            try:
+                if line.startswith(("import ", "import\t")):
+                    # Change by immerrr: don't evaluate import lines to prevent
+                    # code injection into jedi through pth files.
+                    #
+                    # exec(line)
+                    continue
+                line = line.rstrip()
+                dir, dircase = makepath(sitedir, line)
+                if not dircase in known_paths and os.path.exists(dir):
+                    sys_path.append(dir)
+                    known_paths.add(dircase)
+            except Exception:
+                print("Error processing line {:d} of {}:\n".format(n+1, fullname),
+                      file=sys.stderr)
+                import traceback
+                for record in traceback.format_exception(*sys.exc_info()):
+                    for line in record.splitlines():
+                        print('  '+line, file=sys.stderr)
+                print("\nRemainder of file ignored", file=sys.stderr)
+                break
+    if reset:
+        known_paths = None
+    return known_paths
+
+
+def addsitedir(sys_path, sitedir, known_paths=None):
+    """Add 'sitedir' argument to sys_path if missing and handle .pth files in
+    'sitedir'"""
+    if known_paths is None:
+        known_paths = _init_pathinfo(sys_path)
+        reset = 1
+    else:
+        reset = 0
+    sitedir, sitedircase = makepath(sitedir)
+    if not sitedircase in known_paths:
+        sys_path.append(sitedir)        # Add path component
+        known_paths.add(sitedircase)
+    try:
+        names = os.listdir(sitedir)
+    except OSError:
+        return
+    names = [name for name in names if name.endswith(".pth")]
+    for name in sorted(names):
+        addpackage(sys_path, sitedir, name, known_paths)
+    if reset:
+        known_paths = None
+    return known_paths

jedi/evaluate/sys_path.py

@@ -1,7 +1,7 @@
 import glob
 import os
 import sys
-from site import addsitedir
+from jedi.evaluate.site import addsitedir
 
 from jedi._compatibility import exec_function, unicode
 from jedi.parser import tree
@@ -54,12 +54,9 @@ def _get_venv_path_dirs(venv):
     """Get sys.path for venv without starting up the interpreter."""
     venv = os.path.abspath(venv)
     sitedir = _get_venv_sitepackages(venv)
-    sys.path, old_sys_path = [], sys.path
-    try:
-        addsitedir(sitedir)
-        return sys.path
-    finally:
-        sys.path = old_sys_path
+    sys_path = []
+    addsitedir(sys_path, sitedir)
+    return sys_path
 
 
 def _get_venv_sitepackages(venv):

test/test_evaluate/test_sys_path.py

@@ -46,9 +46,14 @@ def test_get_venv_path(venv):
         pjoin('/path', 'from', 'egg-link'),
         pjoin(site_pkgs, '.', 'relative', 'egg-link', 'path'),
         pjoin(site_pkgs, 'dir-from-foo-pth'),
-        pjoin('/path', 'from', 'smth.py'),
-        pjoin('/path', 'from', 'smth.py:extend_path')
     ]
+
+    # Ensure that pth and egg-link paths were added.
+    assert venv_path[:len(ETALON)] == ETALON
+
     # Ensure that none of venv dirs leaked to the interpreter.
     assert not set(sys.path).intersection(ETALON)
-    assert venv_path[:len(ETALON)] == ETALON
+
+    # Ensure that "import ..." lines were ignored.
+    assert pjoin('/path', 'from', 'smth.py') not in venv_path
+    assert pjoin('/path', 'from', 'smth.py:extend_path') not in venv_path