def __init__(self, fn, *args, **kwargs):
if isinstance(fn, type) and issubclass(fn, WithTypeHints):
# Don't treat Fn class objects as callables.
raise ValueError('Use %s() not %s.' % (fn.__name__, fn.__name__))
self.fn = self.make_fn(fn)
# Now that we figure out the label, initialize the super-class.
super(PTransformWithSideInputs, self).__init__()
if (any([isinstance(v, pvalue.PCollection) for v in args]) or
any([isinstance(v, pvalue.PCollection) for v in kwargs.itervalues()])):
raise error.SideInputError(
'PCollection used directly as side input argument. Specify '
'AsIter(pcollection) or AsSingleton(pcollection) to indicate how the '
'PCollection is to be used.')
self.args, self.kwargs, self.side_inputs = util.remove_objects_from_args(
args, kwargs, pvalue.AsSideInput)
self.raw_side_inputs = args, kwargs
# Prevent name collisions with fns of the form '<function <lambda> at ...>'
self._cached_fn = self.fn
# Ensure fn and side inputs are picklable for remote execution.
self.fn = pickler.loads(pickler.dumps(self.fn))
self.args = pickler.loads(pickler.dumps(self.args))
self.kwargs = pickler.loads(pickler.dumps(self.kwargs))
# For type hints, because loads(dumps(class)) != class.
self.fn = self._cached_fn
def __init__(self, fn, *args, **kwargs):
# type: (WithTypeHints, *Any, **Any) -> None
if isinstance(fn, type) and issubclass(fn, WithTypeHints):
# Don't treat Fn class objects as callables.
raise ValueError('Use %s() not %s.' % (fn.__name__, fn.__name__))
self.fn = self.make_fn(fn, bool(args or kwargs))
# Now that we figure out the label, initialize the super-class.
super(PTransformWithSideInputs, self).__init__()
if (any([isinstance(v, pvalue.PCollection) for v in args]) or any(
[isinstance(v, pvalue.PCollection) for v in kwargs.values()])):
raise error.SideInputError(
'PCollection used directly as side input argument. Specify '
'AsIter(pcollection) or AsSingleton(pcollection) to indicate how the '
'PCollection is to be used.')
self.args, self.kwargs, self.side_inputs = util.remove_objects_from_args(
args, kwargs, pvalue.AsSideInput)
self.raw_side_inputs = args, kwargs
# Prevent name collisions with fns of the form '<function <lambda> at ...>'
self._cached_fn = self.fn
# Ensure fn and side inputs are picklable for remote execution.
try:
self.fn = pickler.loads(pickler.dumps(self.fn))
except RuntimeError as e:
raise RuntimeError('Unable to pickle fn %s: %s' % (self.fn, e))
self.args = pickler.loads(pickler.dumps(self.args))
self.kwargs = pickler.loads(pickler.dumps(self.kwargs))
# For type hints, because loads(dumps(class)) != class.
self.fn = self._cached_fn
def __init__(self, fn_or_label, *args, **kwargs):
if fn_or_label is None or isinstance(fn_or_label, basestring):
label = fn_or_label
fn, args = args[0], args[1:]
else:
label = None
fn = fn_or_label
if isinstance(fn, type) and issubclass(fn, typehints.WithTypeHints):
# Don't treat Fn class objects as callables.
raise ValueError('Use %s() not %s.' % (fn.__name__, fn.__name__))
self.fn = self.make_fn(fn)
# Now that we figure out the label, initialize the super-class.
super(PTransformWithSideInputs, self).__init__(label=label)
if (any([isinstance(v, pvalue.PCollection) for v in args]) or any(
[isinstance(v, pvalue.PCollection) for v in kwargs.itervalues()])):
raise error.SideInputError(
'PCollection used directly as side input argument. Specify '
'AsIter(pcollection) or AsSingleton(pcollection) to indicate how the '
'PCollection is to be used.')
self.args, self.kwargs, self.side_inputs = util.remove_objects_from_args(
args, kwargs, pvalue.PCollectionView)
self.raw_side_inputs = args, kwargs
# Prevent name collisions with fns of the form '<function <lambda> at ...>'
self._cached_fn = self.fn
# Ensure fn and side inputs are picklable for remote execution.
self.fn = pickler.loads(pickler.dumps(self.fn))
self.args = pickler.loads(pickler.dumps(self.args))
self.kwargs = pickler.loads(pickler.dumps(self.kwargs))
# For type hints, because loads(dumps(class)) != class.
self.fn = self._cached_fn
def test_remove_objects_from_args_nothing_to_remove(self):
args, kwargs, objs = remove_objects_from_args([1, 2], {
'x': 1,
'y': 2
}, (str, float))
self.assertEqual([1, 2], args)
self.assertEqual({'x': 1, 'y': 2}, kwargs)
self.assertEqual([], objs)
def test_remove_objects_from_args(self):
args, kwargs, objs = remove_objects_from_args([1, 'a'], {
'x': 1,
'y': 3.14
}, (str, float))
self.assertEqual([1, ArgumentPlaceholder()], args)
self.assertEqual({'x': 1, 'y': ArgumentPlaceholder()}, kwargs)
self.assertEqual(['a', 3.14], objs)
def test_remove_objects_from_args_nothing_to_remove(self):
args, kwargs, objs = remove_objects_from_args(
[1, 2], {'x': 1, 'y': 2}, (str, float))
self.assertEquals([1, 2], args)
self.assertEquals({'x': 1, 'y': 2}, kwargs)
self.assertEquals([], objs)
def test_remove_objects_from_args(self):
args, kwargs, objs = remove_objects_from_args(
[1, 'a'], {'x': 1, 'y': 3.14}, (str, float))
self.assertEquals([1, ArgumentPlaceholder()], args)
self.assertEquals({'x': 1, 'y': ArgumentPlaceholder()}, kwargs)
self.assertEquals(['a', 3.14], objs)
