def symmetric_binary_op(state, unused_arg):
# TODO(robertwb): This may not be entirely correct...
b, a = Const.unwrap(state.stack.pop()), Const.unwrap(state.stack.pop())
if a == b:
state.stack.append(a)
elif type(a) == type(b) and isinstance(a, typehints.SequenceTypeConstraint):
state.stack.append(type(a)(union(element_type(a), element_type(b))))
else:
state.stack.append(Any)
def binary_subscr(state, unused_arg):
index = state.stack.pop()
base = Const.unwrap(state.stack.pop())
if base in (str, unicode):
out = base
elif (isinstance(index, Const) and isinstance(index.value, int) and
isinstance(base, typehints.IndexableTypeConstraint)):
try:
out = base._constraint_for_index(index.value)
except IndexError:
out = element_type(base)
elif index == slice and isinstance(base, typehints.IndexableTypeConstraint):
out = base
else:
out = element_type(base)
state.stack.append(out)
def _strip_output_annotations(self, type_hint):
annotations = (TimestampedValue, WindowedValue, pvalue.SideOutputValue)
# TODO(robertwb): These should be parameterized types that the
# type inferencer understands.
if (type_hint in annotations
or trivial_inference.element_type(type_hint) in annotations):
return Any
else:
return type_hint
def default_type_hints(self):
type_hints = get_type_hints(self._fn)
# If the fn was a DoFn annotated with a type-hint that hinted a return
# type compatible with Iterable[Any], then we strip off the outer
# container type due to the 'flatten' portion of FlatMap.
# TODO(robertwb): Should we require an iterable specification for FlatMap?
if type_hints.output_types:
args, kwargs = type_hints.output_types
if len(args) == 1 and is_consistent_with(args[0], Iterable[Any]):
type_hints = type_hints.copy()
type_hints.set_output_types(element_type(args[0]), **kwargs)
return type_hints
def unpack_sequence(state, arg):
t = state.stack.pop()
if isinstance(t, Const):
try:
unpacked = [Const(ti) for ti in t.value]
if len(unpacked) != arg:
unpacked = [Any] * arg
except TypeError:
unpacked = [Any] * arg
elif (isinstance(t, typehints.TupleHint.TupleConstraint)
and len(t.tuple_types) == arg):
unpacked = list(t.tuple_types)
else:
unpacked = [element_type(t)] * arg
state.stack += reversed(unpacked)
def build_map_unpack(state, arg):
"""Joins arg count maps from the stack into a single dict."""
key_types = []
value_types = []
for _ in range(arg):
type_constraint = state.stack.pop()
if isinstance(type_constraint, typehints.Dict.DictConstraint):
key_types.append(type_constraint.key_type)
value_types.append(type_constraint.value_type)
else:
key_type, value_type = key_value_types(
element_type(type_constraint))
key_types.append(key_type)
value_types.append(value_type)
state.stack.append(Dict[Union[key_types], Union[value_types]])
def dict_update(state, arg):
other = state.stack.pop()
base = state.stack[-arg]
if isinstance(base, typehints.Dict.DictConstraint):
base_key_type = base.key_type
base_value_type = base.value_type
else:
base_key_type = Any
base_value_type = Any
if isinstance(other, typehints.Dict.DictConstraint):
other_key_type = other.key_type
other_value_type = other.value_type
else:
other_key_type, other_value_type = key_value_types(element_type(other))
state.stack[-arg] = Dict[union(base_key_type, other_key_type),
union(base_value_type, other_value_type)]
def default_type_hints(self):
fn_hints = get_type_hints(self._fn)
if fn_hints.input_types is None:
return fn_hints
else:
# fn(Iterable[V]) -> V becomes CombineFn(V) -> V
input_args, input_kwargs = fn_hints.input_types
if not input_args:
if len(input_kwargs) == 1:
input_args, input_kwargs = tuple(input_kwargs.values()), {}
else:
raise TypeError('Combiner input type must be specified positionally.')
if not is_consistent_with(input_args[0], Iterable[Any]):
raise TypeCheckError(
'All functions for a Combine PTransform must accept a '
'single argument compatible with: Iterable[Any]. '
'Instead a function with input type: %s was received.'
% input_args[0])
input_args = (element_type(input_args[0]),) + input_args[1:]
# TODO(robertwb): Assert output type is consistent with input type?
hints = fn_hints.copy()
hints.set_input_types(*input_args, **input_kwargs)
return hints
def list_append(state, arg):
new_element_type = Const.unwrap(state.stack.pop())
state.stack[-arg] = List[Union[element_type(state.stack[-arg]),
new_element_type]]
def get_iter(state, unused_arg):
state.stack.append(Iterable[element_type(state.stack.pop())])
def infer_output_type(self, input_type):
return trivial_inference.element_type(
self.fn.infer_output_type(input_type))
def set_update(state, arg):
other = state.stack.pop()
base = state.stack[-arg]
state.stack[-arg] = Set[union(element_type(base), element_type(other))]
def list_extend(state, arg):
tail = state.stack.pop()
base = state.stack[-arg]
state.stack[-arg] = List[union(element_type(base), element_type(tail))]
def list_to_tuple(state, arg):
base = state.stack.pop()
state.stack.append(Tuple[element_type(base), ...])
def set_add(state, arg):
new_element_type = Const.unwrap(state.stack.pop())
state.stack[-arg] = Set[union(element_type(state.stack[-arg]),
new_element_type)]
