[arrow] branch master updated: ARROW-2805: [Python] Use official way to find TensorFlow module

[ro...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

robertnishihara pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/master by this push:
     new 010c874  ARROW-2805: [Python] Use official way to find TensorFlow module
010c874 is described below

commit 010c87402071d715e6fd0c3d22a0b13820b9aed5
Author: Philipp Moritz <pc...@gmail.com>
AuthorDate: Fri Jul 6 20:58:32 2018 -0700

    ARROW-2805: [Python] Use official way to find TensorFlow module
    
    Still trying this out, but this seems like the way to do it.
    
    Note it will use the slow path in Python 2 but that's ok I think.
    
    Author: Philipp Moritz <pc...@gmail.com>
    
    Closes #2224 from pcmoritz/fix-tf-workaround and squashes the following commits:
    
    a874471 <Philipp Moritz> fix module.__path__ usage
    b4ffbd9 <Philipp Moritz> more fixes
    56df127 <Philipp Moritz> fix Python 2.7
    1825372 <Philipp Moritz> use official way to load module
---
 python/pyarrow/compat.py | 48 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/python/pyarrow/compat.py b/python/pyarrow/compat.py
index 1fcaf4c..8be27b7 100644
--- a/python/pyarrow/compat.py
+++ b/python/pyarrow/compat.py
@@ -172,30 +172,50 @@ def import_tensorflow_extension():
     https://github.com/apache/arrow/pull/2096.
     """
     import os
-    import site
     tensorflow_loaded = False
 
     # Try to load the tensorflow extension directly
     # This is a performance optimization, tensorflow will always be
     # loaded via the "import tensorflow" statement below if this
     # doesn't succeed.
+    #
+    # This uses the official way of loading modules from
+    # https://docs.python.org/3/library/importlib.html#approximating-importlib-import-module
+
     try:
-        site_paths = site.getsitepackages() + [site.getusersitepackages()]
-    except AttributeError:
-        # Workaround for https://github.com/pypa/virtualenv/issues/228,
-        # this happends in some configurations of virtualenv
-        site_paths = [os.path.dirname(site.__file__) + '/site-packages']
-    for site_path in site_paths:
-        ext = os.path.join(site_path, "tensorflow",
-                           "libtensorflow_framework.so")
-        if os.path.exists(ext):
-            import ctypes
-            ctypes.CDLL(ext)
-            tensorflow_loaded = True
-            break
+        import importlib
+        absolute_name = importlib.util.resolve_name("tensorflow", None)
+    except (ImportError, AttributeError):
+        # Sometimes, importlib is not available (e.g. Python 2)
+        # or importlib.util is not available (e.g. Python 2.7)
+        spec = None
+    else:
+        import sys
+        for finder in sys.meta_path:
+            try:
+                spec = finder.find_spec(absolute_name, None)
+            except AttributeError:
+                # On Travis (Python 3.5) the above produced:
+                # AttributeError: 'VendorImporter' object has no
+                # attribute 'find_spec'
+                spec = None
+            if spec is not None:
+                break
+
+    if spec:
+        module = importlib.util.module_from_spec(spec)
+        for path in module.__path__:
+            ext = os.path.join(path, "libtensorflow_framework.so")
+            if os.path.exists(ext):
+                import ctypes
+                ctypes.CDLL(ext)
+                tensorflow_loaded = True
+                break
+
 
     # If the above failed, try to load tensorflow the normal way
     # (this is more expensive)
+
     if not tensorflow_loaded:
         try:
             import tensorflow