ENH: derive from C-pickler for fast serialization

.travis.yml

@@ -13,8 +13,8 @@ matrix:
       dist: trusty
       python: "pypy3"
     - os: linux
-      if: commit_message =~ /(\[ci python-nightly\])/
       env: PYTHON_NIGHTLY=1
+      python: 3.7
     - os: linux
       python: 3.7
     - os: linux
@@ -91,8 +91,12 @@ install:
   - $PYTHON_EXE -m pip install .
   - $PYTHON_EXE -m pip install --upgrade -r dev-requirements.txt
   - $PYTHON_EXE -m pip install tornado
-  - if [[ $TRAVIS_PYTHON_VERSION != 'pypy'* && "$PYTHON_NIGHTLY" != 1 ]]; then
-        $PYTHON_EXE -m pip install numpy scipy;
+  - if [[ $TRAVIS_PYTHON_VERSION != 'pypy'* ]]; then
+        if [[ "$PYTHON_NIGHTLY" == "1" ]]; then
+          $PYTHON_EXE -m pip install git+https://github.com/cython/cython git+https://github.com/numpy/numpy;
+        else
+          $PYTHON_EXE -m pip install numpy scipy;
+        fi
     fi
   - if [[ $PROJECT != "" ]]; then
         $PYTHON_EXE -m pip install $TEST_REQUIREMENTS;
@@ -126,5 +130,6 @@ script:
       fi
     fi
 after_success:
+  - pip install coverage codecov
   - coverage combine --append
   - codecov

cloudpickle/__init__.py

@@ -1,5 +1,12 @@
 from __future__ import absolute_import
 
-from cloudpickle.cloudpickle import *
+import sys
+import pickle
+
+
+if sys.version_info[:2] >= (3, 8):
+    from cloudpickle.cloudpickle_fast import *
+else:
+    from cloudpickle.cloudpickle import *
 
 __version__ = '1.2.0.dev0'

cloudpickle/cloudpickle.py

@@ -102,6 +102,8 @@
     PY2 = False
     from importlib._bootstrap import _find_spec
 
+_extract_code_globals_cache = weakref.WeakKeyDictionary()
+
 
 def _ensure_tracking(class_def):
     with _DYNAMIC_CLASS_TRACKER_LOCK:
@@ -195,6 +197,78 @@ def _is_global(obj, name=None):
     return obj2 is obj
 
 
+def _extract_code_globals(co):
+    """
+    Find all globals names read or written to by codeblock co
+    """
+    out_names = _extract_code_globals_cache.get(co)
+    if out_names is None:
+        names = co.co_names
+        out_names = {names[oparg] for _, oparg in _walk_global_ops(co)}
+
+        # Declaring a function inside another one using the "def ..."
+        # syntax generates a constant code object corresonding to the one
+        # of the nested function's As the nested function may itself need
+        # global variables, we need to introspect its code, extract its
+        # globals, (look for code object in it's co_consts attribute..) and
+        # add the result to code_globals
+        if co.co_consts:
+            for const in co.co_consts:
+                if isinstance(const, types.CodeType):
+                    out_names |= _extract_code_globals(const)
+
+        _extract_code_globals_cache[co] = out_names
+
+    return out_names
+
+
+def _find_imported_submodules(code, top_level_dependencies):
+    """
+    Find currently imported submodules used by a function.
+
+    Submodules used by a function need to be detected and referenced for the
+    function to work correctly at depickling time. Because submodules can be
+    referenced as attribute of their parent package (``package.submodule``), we
+    need a special introspection technique that does not rely on GLOBAL-related
+    opcodes to find references of them in a code object.
+
+    Example:
+    ```
+    import concurrent.futures
+    import cloudpickle
+    def func():
+        x = concurrent.futures.ThreadPoolExecutor
+    if __name__ == '__main__':
+        cloudpickle.dumps(func)
+    ```
+    The globals extracted by cloudpickle in the function's state include the
+    concurrent package, but not its submodule (here, concurrent.futures), which
+    is the module used by func. Find_imported_submodules will detect the usage
+    of concurrent.futures. Saving this module alongside with func will ensure
+    that calling func once depickled does not fail due to concurrent.futures
+    not being imported
+    """
+
+    subimports = []
+    # check if any known dependency is an imported package
+    for x in top_level_dependencies:
+        if (isinstance(x, types.ModuleType) and
+                hasattr(x, '__package__') and x.__package__):
+            # check if the package has any currently loaded sub-imports
+            prefix = x.__name__ + '.'
+            # A concurrent thread could mutate sys.modules,
+            # make sure we iterate over a copy to avoid exceptions
+            for name in list(sys.modules):
+                # Older versions of pytest will add a "None" module to
+                # sys.modules.
+                if name is not None and name.startswith(prefix):
+                    # check whether the function can address the sub-module
+                    tokens = set(name[len(prefix):].split('.'))
+                    if not tokens - set(code.co_names):
+                        subimports.append(sys.modules[name])
+    return subimports
+
+
 def _make_cell_set_template_code():
     """Get the Python compiler to emit LOAD_FAST(arg); STORE_DEREF
 
@@ -676,10 +750,16 @@ def save_function_tuple(self, func):
         save(_fill_function)  # skeleton function updater
         write(pickle.MARK)    # beginning of tuple that _fill_function expects
 
-        self._save_subimports(
+        submodules = _find_imported_submodules(
             code,
             itertools.chain(f_globals.values(), closure_values or ()),
         )
+        for s in submodules:
+            # ensure that subimport s is loaded at unpickling time
+            self.save(s)
+            # then discards the reference to it
+            self.write(pickle.POP)
+
 
         # create a skeleton function object and memoize it
         save(_make_skel_func)
@@ -711,28 +791,6 @@ def save_function_tuple(self, func):
         write(pickle.TUPLE)
         write(pickle.REDUCE)  # applies _fill_function on the tuple
 
-    _extract_code_globals_cache = weakref.WeakKeyDictionary()
-
-    @classmethod
-    def extract_code_globals(cls, co):
-        """
-        Find all globals names read or written to by codeblock co
-        """
-        out_names = cls._extract_code_globals_cache.get(co)
-        if out_names is None:
-            names = co.co_names
-            out_names = {names[oparg] for _, oparg in _walk_global_ops(co)}
-
-            # see if nested function have any global refs
-            if co.co_consts:
-                for const in co.co_consts:
-                    if isinstance(const, types.CodeType):
-                        out_names |= cls.extract_code_globals(const)
-
-            cls._extract_code_globals_cache[co] = out_names
-
-        return out_names
-
     def extract_func_data(self, func):
         """
         Turn the function into a tuple of data necessary to recreate it:
@@ -741,7 +799,7 @@ def extract_func_data(self, func):
         code = func.__code__
 
         # extract all global ref's
-        func_global_refs = self.extract_code_globals(code)
+        func_global_refs = _extract_code_globals(code)
 
         # process all variables referenced by global environment
         f_globals = {}