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 __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 __init__(self, fn_or_label, *args, **kwargs):
if fn_or_label is None:
fn_or_label = str(args[0])
if isinstance(fn_or_label, str):
label = fn_or_label
self.fn, args = self.make_fn(args[0]), args[1:]
else:
label = None
self.fn = self.make_fn(fn_or_label)
# 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.AsSingleton, pvalue.AsIter))
self.raw_side_inputs = args, kwargs
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)