FIX ApiModule can be vendored in the pytest package (cloudpipe#135)

HyukjinKwon · Feb 5, 2018 · 22af562 · 22af562
1 parent dac351a
commit 22af562
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py
@@ -690,10 +690,10 @@ def f():
             s = py.builtin.set([1])
             return s.pop()
 
-        # some setup is required to allow pytest apimodules to be correctly serializable.
+        # some setup is required to allow pytest apimodules to be correctly
+        # serializable.
         from cloudpickle import CloudPickler
-        from py._apipkg import ApiModule
-        CloudPickler.dispatch[ApiModule] = CloudPickler.save_module
+        CloudPickler.dispatch[type(py.builtin)] = CloudPickler.save_module
         g = cloudpickle.loads(cloudpickle.dumps(f))
 
         result = g()