#

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

#

"""

Unit tests for typecheck.

See additional runtime_type_check=True tests in ptransform_test.py.

"""

# pytype: skip-file

import os

import tempfile

import unittest

from typing import Iterable

from typing import Tuple

import apache_beam as beam

from apache_beam import Pipeline

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.options.pipeline_options import TypeOptions

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.typehints import decorators

from apache_beam.typehints import with_input_types

from apache_beam.typehints import with_output_types

decorators._enable_from_callable = True

# Disable frequent lint warning due to pipe operator for chaining transforms.

# pylint: disable=expression-not-assigned

class MyDoFn(beam.DoFn):

def __init__(self, output_filename):

super().__init__()

self.output_filename = output_filename

def _output(self):

"""Returns a file used to record function calls."""

if not hasattr(self, 'output_file'):

self._output_file = open(self.output_filename, 'at', buffering=1)

return self._output_file

def start_bundle(self):

self._output().write('start_bundle\n')

def finish_bundle(self):

self._output().write('finish_bundle\n')

def setup(self):

self._output().write('setup\n')

def teardown(self):

self._output().write('teardown\n')

self._output().close()

def process(self, element: int, *args, **kwargs) -> Iterable[int]:

self._output().write('process\n')

yield element

class MyDoFnBadAnnotation(MyDoFn):

def process(self, element: int, *args, **kwargs) -> int:

# Should raise an exception about return type not being iterable.

return super().process()

class RuntimeTypeCheckTest(unittest.TestCase):

def setUp(self):

# Use FnApiRunner since it guarantees all lifecycle methods will be called.

self.p = TestPipeline(

'FnApiRunner',

options=PipelineOptions(

runtime_type_check=True, performance_runtime_type_check=False))

def test_setup(self):

# Verifies that runtime type checking is enabled for test cases.

def fn(e: int) -> int:

return str(e) # type: ignore

with self.assertRaisesRegex(Exception,

r'output should be.*int.*received.*str'):

_ = self.p | beam.Create([1, 2, 3]) | beam.Map(fn)

self.p.run()

def test_wrapper_pass_through(self):

# We use a file to check the result because the MyDoFn instance passed is

# not the same one that actually runs in the pipeline (it is serialized

# here and deserialized in the worker).

with tempfile.TemporaryDirectory() as tmp_dirname:

path = os.path.join(tmp_dirname + "tmp_filename")

dofn = MyDoFn(path)

result = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn)

assert_that(result, equal_to([1, 2, 3]))

self.p.run()

with open(path, mode="r") as ft:

lines = [line.strip() for line in ft]

self.assertListEqual([

'setup',

'start_bundle',

'process',

'process',

'process',

'finish_bundle',

'teardown',

],

lines)

def test_wrapper_pipeline_type_check(self):

# Verifies that type hints are not masked by the wrapper. What actually

# happens is that the wrapper is applied during self.p.run() (not invoked

# in this case), while pipeline type checks happen during pipeline creation.

# Thus, the wrapper does not have to implement: default_type_hints,

# infer_output_type, get_type_hints.

with tempfile.NamedTemporaryFile(mode='w+t') as f:

dofn = MyDoFnBadAnnotation(f.name)

with self.assertRaisesRegex(ValueError, r'int.*is not iterable'):

_ = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn)

class PerformanceRuntimeTypeCheckTest(unittest.TestCase):

def setUp(self):

self.p = Pipeline(

options=PipelineOptions(

performance_runtime_type_check=True, pipeline_type_check=False))

def assertStartswith(self, msg, prefix):

self.assertTrue(

msg.startswith(prefix), '"%s" does not start with "%s"' % (msg, prefix))

def test_simple_input_error(self):

with self.assertRaisesRegex(Exception,

"Type-hint for argument: 'x' violated. "

"Expected an instance of {}, "

"instead found 1, an instance of {}".format(

str, int)):

(

self.p

| beam.Create([1, 1])

| beam.FlatMap(lambda x: [int(x)]).with_input_types(

str).with_output_types(int))

self.p.run().wait_until_finish()

def test_simple_output_error(self):

with self.assertRaisesRegex(Exception,

"Type-hint for argument: 'x' violated. "):

(

self.p

| beam.Create(['1', '1'])

| beam.FlatMap(lambda x: [int(x)]).with_input_types(

int).with_output_types(int))

self.p.run().wait_until_finish()

def test_simple_input_error_with_kwarg_typehints(self):

@with_input_types(element=int)

@with_output_types(int)

class ToInt(beam.DoFn):

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

yield int(element)

with self.assertRaisesRegex(Exception,

"Type-hint for argument: 'element' violated"):

(self.p | beam.Create(['1', '1']) | beam.ParDo(ToInt()))

self.p.run().wait_until_finish()

def test_bad_flatten_input(self):

with self.assertRaisesRegex(

TypeError,

"Inputs to Flatten cannot include an iterable of PCollections. "):

with beam.Pipeline() as p:

pc = p | beam.Create([1, 1])

flatten_inputs = [pc, (pc, )]

flatten_inputs | beam.Flatten()

def test_do_fn_returning_non_iterable_throws_error(self):

# This function is incorrect because it returns a non-iterable object

def incorrect_par_do_fn(x):

return x + 5

with self.assertRaisesRegex(Exception, "'int' object is not iterable "):

(self.p | beam.Create([1, 1]) | beam.FlatMap(incorrect_par_do_fn))

self.p.run().wait_until_finish()

def test_simple_type_satisfied(self):

@with_input_types(int, int)

@with_output_types(int)

class AddWithNum(beam.DoFn):

def process(self, element, num):

return [element + num]

results = (

self.p

| 'T' >> beam.Create([1, 2, 3]).with_output_types(int)

| 'Add' >> beam.ParDo(AddWithNum(), 1))

assert_that(results, equal_to([2, 3, 4]))

self.p.run()

def test_simple_type_violation(self):

self.p._options.view_as(TypeOptions).pipeline_type_check = False

@with_output_types(str)

@with_input_types(x=int)

def int_to_string(x):

return str(x)

(

self.p

| 'Create' >> beam.Create(['some_string'])

| 'ToStr' >> beam.Map(int_to_string))

with self.assertRaisesRegex(Exception,

"Type-hint for argument: 'x' violated. "):

self.p.run().wait_until_finish()

def test_pipeline_checking_satisfied_but_run_time_types_violate(self):

self.p._options.view_as(TypeOptions).pipeline_type_check = False

@with_output_types(Tuple[bool, int])

@with_input_types(a=int)

def is_even_as_key(a):

# Simulate a programming error, should be: return (a % 2 == 0, a)

# However this returns Tuple[int, int]

return (a % 2, a)

(

self.p

| 'Nums' >> beam.Create(range(1)).with_output_types(int)

| 'IsEven' >> beam.Map(is_even_as_key)

| 'Parity' >> beam.GroupByKey())

with self.assertRaisesRegex(

Exception,

("Type-hint for return type violated: Tuple\\[<class 'bool'>, <class " +

"'int'>\\] hint type-constraint violated. The type of element #0 in " +

"the passed tuple is incorrect. Expected an instance of type <class " +

"'bool'>, instead received an instance of type int.")):

self.p.run().wait_until_finish()

def test_pipeline_runtime_checking_violation_composite_type_output(self):

self.p._options.view_as(TypeOptions).pipeline_type_check = False

# The type-hinted applied via the 'returns()' method indicates the ParDo

# should return an instance of type: Tuple[float, int]. However, an instance

# of 'int' will be generated instead.

with self.assertRaisesRegex(

Exception,

("Type-hint for return type violated. Expected an instance of {}, " +

"instead found 4.0, an instance of {}.").format(int, float)):

(

self.p

| beam.Create([(1, 3.0)])

| (

'Swap' >>

beam.FlatMap(lambda x_y1: [x_y1[0] + x_y1[1]]).with_input_types(

Tuple[int, float]).with_output_types(int)))

self.p.run().wait_until_finish()

def test_downstream_input_type_hint_error_has_descriptive_error_msg(self):

@with_input_types(int)

@with_output_types(int)

class IntToInt(beam.DoFn):

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

yield element

@with_input_types(str)

@with_output_types(int)

class StrToInt(beam.DoFn):

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

yield int(element)

# This will raise a type check error in IntToInt even though the actual

# type check error won't happen until StrToInt. The user will be told that

# StrToInt's input type hints were not satisfied while running IntToInt.

with self.assertRaisesRegex(

Exception,

("Type-hint for argument: 'element' violated. Expected an instance of "

+ "{}, instead found 9, an instance of {}.").format(str, int)):

(

self.p

| beam.Create([9])

| beam.ParDo(IntToInt())

| beam.ParDo(StrToInt()))

self.p.run().wait_until_finish()

if __name__ == '__main__':

unittest.main()