Source code for apache_beam.typehints.typecheck

## Licensed to the Apache Software Foundation (ASF) under one or more# contributor license agreements. See the NOTICE file distributed with# this work for additional information regarding copyright ownership.# The ASF licenses this file to You under the Apache License, Version 2.0# (the "License"); you may not use this file except in compliance with# the License. You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.#"""Runtime type checking support.For internal use only; no backwards-compatibility guarantees."""importcollectionsimportinspectimportsysimporttypesfromapache_beam.pvalueimportTaggedOutputfromapache_beam.transforms.coreimportDoFnfromapache_beam.transforms.windowimportWindowedValuefromapache_beam.typehints.decoratorsimportGeneratorWrapperfromapache_beam.typehints.decoratorsimportTypeCheckErrorfromapache_beam.typehints.decoratorsimport_check_instance_typefromapache_beam.typehints.decoratorsimportgetcallargs_forhintsfromapache_beam.typehints.typehintsimportCompositeTypeHintErrorfromapache_beam.typehints.typehintsimportSimpleTypeHintErrorfromapache_beam.typehints.typehintsimportcheck_constraint

[docs]classOutputCheckWrapperDoFn(AbstractDoFnWrapper):"""A DoFn that verifies against common errors in the output type."""def__init__(self,dofn,full_label):super(OutputCheckWrapperDoFn,self).__init__(dofn)self.full_label=full_label

def_check_type(self,output):ifoutputisNone:returnoutputelifisinstance(output,(dict,basestring)):object_type=type(output).__name__raiseTypeCheckError('Returning a %s from a ParDo or FlatMap is ''discouraged. Please use list("%s") if you really ''want this behavior.'%(object_type,output))elifnotisinstance(output,collections.Iterable):raiseTypeCheckError('FlatMap and ParDo must return an ''iterable. %s was returned instead.'%type(output))returnoutput

[docs]classTypeCheckWrapperDoFn(AbstractDoFnWrapper):"""A wrapper around a DoFn which performs type-checking of input and output. """def__init__(self,dofn,type_hints,label=None):super(TypeCheckWrapperDoFn,self).__init__(dofn)self.dofn=dofnself._process_fn=self.dofn._process_argspec_fn()iftype_hints.input_types:input_args,input_kwargs=type_hints.input_typesself._input_hints=getcallargs_forhints(self._process_fn,*input_args,**input_kwargs)else:self._input_hints=None# TODO(robertwb): Multi-output.self._output_type_hint=type_hints.simple_output_type(label)

def_type_check_result(self,transform_results):ifself._output_type_hintisNoneortransform_resultsisNone:returntransform_resultsdeftype_check_output(o):# TODO(robertwb): Multi-output.x=o.valueifisinstance(o,(TaggedOutput,WindowedValue))elseoself._type_check(self._output_type_hint,x,is_input=False)# If the return type is a generator, then we will need to interleave our# type-checking with its normal iteration so we don't deplete the# generator initially just by type-checking its yielded contents.ifisinstance(transform_results,types.GeneratorType):returnGeneratorWrapper(transform_results,type_check_output)forointransform_results:type_check_output(o)returntransform_resultsdef_type_check(self,type_constraint,datum,is_input):"""Typecheck a PTransform related datum according to a type constraint. This function is used to optionally type-check either an input or an output to a PTransform. Args: type_constraint: An instance of a typehints.TypeContraint, one of the white-listed builtin Python types, or a custom user class. datum: An instance of a Python object. is_input: True if 'datum' is an input to a PTransform's DoFn. False otherwise. Raises: TypeError: If 'datum' fails to type-check according to 'type_constraint'. """datum_type='input'ifis_inputelse'output'try:check_constraint(type_constraint,datum)exceptCompositeTypeHintErrorase:raiseTypeCheckError,e.message,sys.exc_info()[2]exceptSimpleTypeHintError:error_msg=("According to type-hint expected %s should be of type %s. ""Instead, received '%s', an instance of type %s."%(datum_type,type_constraint,datum,type(datum)))raiseTypeCheckError,error_msg,sys.exc_info()[2]