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)]