Eliminate the fallback to dill when pickling fails for data

If desired, a DillPickler can be explicitly used instead.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=117143518

Author: robertwb

google/cloud/dataflow/coders/coders.py

@@ -236,7 +236,7 @@ def is_deterministic(self):
     return True
 
 
-def dumps(o):
+def maybe_dill_dumps(o):
   """Pickle using cPickle or the Dill pickler as a fallback."""
   # We need to use the dill pickler for objects of certain custom classes,
   # including, for example, ones that contain lambdas.
@@ -246,19 +246,16 @@ def dumps(o):
     return dill.dumps(o)
 
 
-def loads(o):
+def maybe_dill_loads(o):
   """Unpickle using cPickle or the Dill pickler as a fallback."""
   try:
     return pickle.loads(o)
   except Exception:  # pylint: disable=broad-except
     return dill.loads(o)
 
 
-class PickleCoder(FastCoder):
-  """Coder using Python's pickle functionality."""
-
-  def _create_impl(self):
-    return coder_impl.CallbackCoderImpl(dumps, loads)
+class _PickleCoderBase(FastCoder):
+  """Base class for pickling coders."""
 
   def is_deterministic(self):
     # Note that the default coder, the PickleCoder, is not deterministic (for
@@ -268,7 +265,7 @@ def is_deterministic(self):
     return False
 
   def as_cloud_object(self, is_pair_like=True):
-    value = super(PickleCoder, self).as_cloud_object()
+    value = super(_PickleCoderBase, self).as_cloud_object()
     # We currently use this coder in places where we cannot infer the coder to
     # use for the value type in a more granular way.  In places where the
     # service expects a pair, it checks for the "is_pair_like" key, in which
@@ -295,6 +292,20 @@ def value_coder(self):
     return self
 
 
+class PickleCoder(_PickleCoderBase):
+  """Coder using Python's pickle functionality."""
+
+  def _create_impl(self):
+    return coder_impl.CallbackCoderImpl(pickle.dumps, pickle.loads)
+
+
+class DillCoder(_PickleCoderBase):
+  """Coder using dill's pickle functionality."""
+
+  def _create_impl(self):
+    return coder_impl.CallbackCoderImpl(maybe_dill_dumps, maybe_dill_loads)
+
+
 class DeterministicPickleCoder(FastCoder):
   """Throws runtime errors when pickling non-deterministic values."""
 

google/cloud/dataflow/coders/coders_test_common.py

@@ -19,9 +19,21 @@
 import sys
 import unittest
 
+import dill
+
 import coders
 
 
+# Defined out of line for picklability.
+class CustomCoder(coders.Coder):
+
+  def encode(self, x):
+    return str(x+1)
+
+  def decode(self, encoded):
+    return int(encoded) - 1
+
+
 class CodersTest(unittest.TestCase):
 
   # These class methods ensure that we test each defined coder in both
@@ -34,9 +46,10 @@ def setUpClass(cls):
 
   @classmethod
   def tearDownClass(cls):
-    standard = set(
-        c for c in coders.__dict__.values()
-        if isinstance(c, type) and issubclass(c, coders.Coder))
+    standard = set(c
+                   for c in coders.__dict__.values()
+                   if isinstance(c, type) and issubclass(c, coders.Coder) and
+                   'Base' not in c.__name__)
     standard -= set([coders.Coder, coders.ToStringCoder, coders.FloatCoder,
                      coders.Base64PickleCoder, coders.FastCoder,
                      coders.WindowCoder, coders.WindowedValueCoder])
@@ -59,15 +72,12 @@ def check_coder(self, coder, *values):
     self._observe(coder)
     for v in values:
       self.assertEqual(v, coder.decode(coder.encode(v)))
+    copy1 = dill.loads(dill.dumps(coder))
+    copy2 = dill.loads(dill.dumps(coder))
+    for v in values:
+      self.assertEqual(v, copy1.decode(copy2.encode(v)))
 
   def test_custom_coder(self):
-    class CustomCoder(coders.Coder):
-
-      def encode(self, x):
-        return str(x+1)
-
-      def decode(self, encoded):
-        return int(encoded) - 1
 
     self.check_coder(CustomCoder(), 1, -10, 5)
     self.check_coder(coders.TupleCoder((CustomCoder(), coders.BytesCoder())),
@@ -87,6 +97,13 @@ def test_deterministic_pickle_coder(self):
     self.check_coder(coders.TupleCoder((coder, coders.PickleCoder())),
                      (1, dict()), ('a', [dict()]))
 
+  def test_dill_coder(self):
+    cell_value = (lambda x: lambda: x)(0).func_closure[0]
+    self.check_coder(coders.DillCoder(), 'a', 1, cell_value)
+    self.check_coder(
+        coders.TupleCoder((coders.VarIntCoder(), coders.DillCoder())),
+        (1, cell_value))
+
   def test_bytes_coder(self):
     self.check_coder(coders.BytesCoder(), 'a', '\0', 'z' * 1000)