#

# 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.

"""

# pytype: skip-file

import inspect

import sys

import types

from collections import abc

from apache_beam import pipeline

from apache_beam.pvalue import TaggedOutput

from apache_beam.transforms import core

from apache_beam.transforms.core import DoFn

from apache_beam.transforms.window import TimestampedValue

from apache_beam.transforms.window import WindowedValue

from apache_beam.typehints.decorators import GeneratorWrapper

from apache_beam.typehints.decorators import TypeCheckError

from apache_beam.typehints.decorators import _check_instance_type

from apache_beam.typehints.decorators import getcallargs_forhints

from apache_beam.typehints.typehints import CompositeTypeHintError

from apache_beam.typehints.typehints import SimpleTypeHintError

from apache_beam.typehints.typehints import check_constraint

from apache_beam.typehints.typehints import normalize

class AbstractDoFnWrapper(DoFn):

"""An abstract class to create wrapper around DoFn"""

def __init__(self, dofn):

super().__init__()

self.dofn = dofn

def __getattribute__(self, name):

if (name.startswith('_') or name in self.__dict__ or

hasattr(type(self), name)):

return object.__getattribute__(self, name)

else:

return getattr(self.dofn, name)

def _inspect_start_bundle(self):

return self.dofn.get_function_arguments('start_bundle')

def _inspect_process(self):

return self.dofn.get_function_arguments('process')

def _inspect_finish_bundle(self):

return self.dofn.get_function_arguments('finish_bundle')

def wrapper(self, method, args, kwargs):

return method(*args, **kwargs)

def setup(self):

return self.dofn.setup()

def start_bundle(self, *args, **kwargs):

return self.wrapper(self.dofn.start_bundle, args, kwargs)

def process(self, *args, **kwargs):

return self.wrapper(self.dofn.process, args, kwargs)

def finish_bundle(self, *args, **kwargs):

return self.wrapper(self.dofn.finish_bundle, args, kwargs)

def teardown(self):

return self.dofn.teardown()

class OutputCheckWrapperDoFn(AbstractDoFnWrapper):

"""A DoFn that verifies against common errors in the output type."""

def __init__(self, dofn, full_label):

super().__init__(dofn)

self.full_label = full_label

def wrapper(self, method, args, kwargs):

try:

result = method(*args, **kwargs)

except TypeCheckError as e:

# TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name

error_msg = (

'Runtime type violation detected within ParDo(%s): '

'%s' % (self.full_label, e))

_, _, tb = sys.exc_info()

raise TypeCheckError(error_msg).with_traceback(tb)

else:

return self._check_type(result)

@staticmethod

def _check_type(output):

if output is None:

return output

elif isinstance(output, (dict, bytes, str)):

object_type = type(output).__name__

raise TypeCheckError(

'Returning a %s from a ParDo or FlatMap is '

'discouraged. Please use list("%s") if you really '

'want this behavior.' % (object_type, output))

elif not isinstance(output, abc.Iterable):

raise TypeCheckError(

'FlatMap and ParDo must return an '

'iterable. %s was returned instead.' % type(output))

return output

class TypeCheckWrapperDoFn(AbstractDoFnWrapper):

"""A wrapper around a DoFn which performs type-checking of input and output.

"""

def __init__(self, dofn, type_hints, label=None):

super().__init__(dofn)

self._process_fn = self.dofn._process_argspec_fn()

if type_hints.input_types:

input_args, input_kwargs = type_hints.input_types

self._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 wrapper(self, method, args, kwargs):

result = method(*args, **kwargs)

return self._type_check_result(result)

def process(self, *args, **kwargs):

if self._input_hints:

actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs) # pylint: disable=deprecated-method

for var, hint in self._input_hints.items():

if hint is actual_inputs[var]:

# self parameter

continue

_check_instance_type(hint, actual_inputs[var], var, True)

return self._type_check_result(self.dofn.process(*args, **kwargs))

def _type_check_result(self, transform_results):

if self._output_type_hint is None or transform_results is None:

return transform_results

def type_check_output(o):

if isinstance(o, TimestampedValue) and hasattr(o, "__orig_class__"):

# when a typed TimestampedValue is set, check the value type

x = o.value

# per https://stackoverflow.com/questions/57706180/,

# __orig_class__ is te safe way to obtain the actual type

# from from Generic[T], supported since Python 3.5.3

beam_type = normalize(o.__orig_class__.__args__[0])

self.type_check(beam_type, x, is_input=False)

else:

# TODO(robertwb): Multi-output.

x = o.value if isinstance(o, (TaggedOutput, WindowedValue)) else o

self.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.

if isinstance(transform_results, types.GeneratorType):

return GeneratorWrapper(transform_results, type_check_output)

for o in transform_results:

type_check_output(o)

return transform_results

@staticmethod

def type_check(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' if is_input else 'output'

try:

check_constraint(type_constraint, datum)

except CompositeTypeHintError as e:

_, _, tb = sys.exc_info()

raise TypeCheckError(e.args[0]).with_traceback(tb)

except SimpleTypeHintError:

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

_, _, tb = sys.exc_info()

raise TypeCheckError(error_msg).with_traceback(tb)

class TypeCheckCombineFn(core.CombineFn):

"""A wrapper around a CombineFn performing type-checking of input and output.

"""

def __init__(self, combinefn, type_hints, label=None):

self._combinefn = combinefn

self._input_type_hint = type_hints.input_types

self._output_type_hint = type_hints.simple_output_type(label)

self._label = label

def setup(self, *args, **kwargs):

self._combinefn.setup(*args, **kwargs)

def create_accumulator(self, *args, **kwargs):

return self._combinefn.create_accumulator(*args, **kwargs)

def add_input(self, accumulator, element, *args, **kwargs):

if self._input_type_hint:

try:

_check_instance_type(

self._input_type_hint[0][0].tuple_types[1],

element,

'element',

True)

except TypeCheckError as e:

error_msg = (

'Runtime type violation detected within %s: '

'%s' % (self._label, e))

_, _, tb = sys.exc_info()

raise TypeCheckError(error_msg).with_traceback(tb)

return self._combinefn.add_input(accumulator, element, *args, **kwargs)

def merge_accumulators(self, accumulators, *args, **kwargs):

return self._combinefn.merge_accumulators(accumulators, *args, **kwargs)

def compact(self, accumulator, *args, **kwargs):

return self._combinefn.compact(accumulator, *args, **kwargs)

def extract_output(self, accumulator, *args, **kwargs):

result = self._combinefn.extract_output(accumulator, *args, **kwargs)

if self._output_type_hint:

try:

_check_instance_type(

self._output_type_hint.tuple_types[1], result, None, True)

except TypeCheckError as e:

error_msg = (

'Runtime type violation detected within %s: '

'%s' % (self._label, e))

_, _, tb = sys.exc_info()

raise TypeCheckError(error_msg).with_traceback(tb)

return result

def teardown(self, *args, **kwargs):

self._combinefn.teardown(*args, **kwargs)

class TypeCheckVisitor(pipeline.PipelineVisitor):

_in_combine = False

def enter_composite_transform(self, applied_transform):

if isinstance(applied_transform.transform, core.CombinePerKey):

self._in_combine = True

self._wrapped_fn = applied_transform.transform.fn = TypeCheckCombineFn(

applied_transform.transform.fn,

applied_transform.transform.get_type_hints(),

applied_transform.full_label)

def leave_composite_transform(self, applied_transform):

if isinstance(applied_transform.transform, core.CombinePerKey):

self._in_combine = False

def visit_transform(self, applied_transform):

transform = applied_transform.transform

if isinstance(transform, core.ParDo):

if self._in_combine:

if isinstance(transform.fn, core.CombineValuesDoFn):

transform.fn.combinefn = self._wrapped_fn

else:

transform.fn = transform.dofn = OutputCheckWrapperDoFn(

TypeCheckWrapperDoFn(

transform.fn,

transform.get_type_hints(),

applied_transform.full_label),

applied_transform.full_label)

class PerformanceTypeCheckVisitor(pipeline.PipelineVisitor):

def visit_transform(self, applied_transform):

transform = applied_transform.transform

full_label = applied_transform.full_label

# Store output type hints in current transform

output_type_hints = self.get_output_type_hints(transform)

if output_type_hints:

transform._add_type_constraint_from_consumer(

full_label, output_type_hints)

# Store input type hints in producer transform

input_type_hints = self.get_input_type_hints(transform)

if input_type_hints and len(applied_transform.inputs):

producer = applied_transform.inputs[0].producer

if producer:

producer.transform._add_type_constraint_from_consumer(

full_label, input_type_hints)

def get_input_type_hints(self, transform):

type_hints = transform.get_type_hints()

input_types = None

if type_hints.input_types:

normal_hints, kwarg_hints = type_hints.input_types

if kwarg_hints:

input_types = kwarg_hints

if normal_hints:

input_types = normal_hints

parameter_name = 'Unknown Parameter'

if hasattr(transform, 'fn'):

try:

argspec = inspect.getfullargspec(transform.fn._process_argspec_fn())

except TypeError:

# An unsupported callable was passed to getfullargspec

pass

else:

if len(argspec.args):

arg_index = 0

if argspec.args[0] == 'self' and len(argspec.args) > 1:

arg_index = 1

parameter_name = argspec.args[arg_index]

if isinstance(input_types, dict):

input_types = (input_types[argspec.args[arg_index]], )

if input_types and len(input_types):

input_types = input_types[0]

return parameter_name, input_types

def get_output_type_hints(self, transform):

type_hints = transform.get_type_hints()

output_types = None

if type_hints.output_types:

normal_hints, kwarg_hints = type_hints.output_types

if kwarg_hints:

output_types = kwarg_hints

if normal_hints:

output_types = normal_hints

if output_types and len(output_types):

output_types = output_types[0]

return None, output_types