#
# 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 the PTransform and descendants."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import operator
import re
import sys
import typing
import unittest
from builtins import map
from builtins import range
from builtins import zip
from functools import reduce

# patches unittest.TestCase to be python3 compatible
import future.tests.base  # pylint: disable=unused-import
import hamcrest as hc
from nose.plugins.attrib import attr

import apache_beam as beam
import apache_beam.pvalue as pvalue
import apache_beam.transforms.combiners as combine
import apache_beam.typehints as typehints
from apache_beam.io.iobase import Read
from apache_beam.metrics import Metrics
from apache_beam.metrics.metric import MetricsFilter
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.transforms import window
from apache_beam.transforms.core import _GroupByKeyOnly
from apache_beam.transforms.display import DisplayData
from apache_beam.transforms.display import DisplayDataItem
from apache_beam.transforms.ptransform import PTransform
from apache_beam.typehints import with_input_types
from apache_beam.typehints import with_output_types
from apache_beam.typehints.typehints_test import TypeHintTestCase
from apache_beam.utils.windowed_value import WindowedValue

# Disable frequent lint warning due to pipe operator for chaining transforms.
# pylint: disable=expression-not-assigned


class PTransformTest(unittest.TestCase):
  # Enable nose tests running in parallel
  _multiprocess_can_split_ = True

  @classmethod
  def setUpClass(cls):
    # Method has been renamed in Python 3
    if sys.version_info[0] < 3:
      cls.assertCountEqual = cls.assertItemsEqual

  def assertStartswith(self, msg, prefix):
    self.assertTrue(msg.startswith(prefix),
                    '"%s" does not start with "%s"' % (msg, prefix))

  def test_str(self):
    self.assertEqual('<PTransform(PTransform) label=[PTransform]>',
                     str(PTransform()))

    pa = TestPipeline()
    res = pa | 'ALabel' >> beam.Create([1, 2])
    self.assertEqual('AppliedPTransform(ALabel/Read, Read)',
                     str(res.producer))

    pc = TestPipeline()
    res = pc | beam.Create([1, 2])
    inputs_tr = res.producer.transform
    inputs_tr.inputs = ('ci',)
    self.assertEqual(
        """<Read(PTransform) label=[Read] inputs=('ci',)>""",
        str(inputs_tr))

    pd = TestPipeline()
    res = pd | beam.Create([1, 2])
    side_tr = res.producer.transform
    side_tr.side_inputs = (4,)
    self.assertEqual(
        '<Read(PTransform) label=[Read] side_inputs=(4,)>',
        str(side_tr))

    inputs_tr.side_inputs = ('cs',)
    self.assertEqual(
        """<Read(PTransform) label=[Read] """
        """inputs=('ci',) side_inputs=('cs',)>""",
        str(inputs_tr))

  def test_do_with_do_fn(self):
    class AddNDoFn(beam.DoFn):

      def process(self, element, addon):
        return [element + addon]

    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    result = pcoll | 'Do' >> beam.ParDo(AddNDoFn(), 10)
    assert_that(result, equal_to([11, 12, 13]))
    pipeline.run()

  def test_do_with_unconstructed_do_fn(self):
    class MyDoFn(beam.DoFn):

      def process(self):
        pass

    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    with self.assertRaises(ValueError):
      pcoll | 'Do' >> beam.ParDo(MyDoFn)  # Note the lack of ()'s

  def test_do_with_callable(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    result = pcoll | 'Do' >> beam.FlatMap(lambda x, addon: [x + addon], 10)
    assert_that(result, equal_to([11, 12, 13]))
    pipeline.run()

  def test_do_with_side_input_as_arg(self):
    pipeline = TestPipeline()
    side = pipeline | 'Side' >> beam.Create([10])
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    result = pcoll | 'Do' >> beam.FlatMap(
        lambda x, addon: [x + addon], pvalue.AsSingleton(side))
    assert_that(result, equal_to([11, 12, 13]))
    pipeline.run()

  def test_do_with_side_input_as_keyword_arg(self):
    pipeline = TestPipeline()
    side = pipeline | 'Side' >> beam.Create([10])
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    result = pcoll | 'Do' >> beam.FlatMap(
        lambda x, addon: [x + addon], addon=pvalue.AsSingleton(side))
    assert_that(result, equal_to([11, 12, 13]))
    pipeline.run()

  def test_do_with_do_fn_returning_string_raises_warning(self):
    pipeline = TestPipeline()
    pipeline._options.view_as(TypeOptions).runtime_type_check = True
    pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])
    pcoll | 'Do' >> beam.FlatMap(lambda x: x + '1')

    # Since the DoFn directly returns a string we should get an error warning
    # us.
    with self.assertRaises(typehints.TypeCheckError) as cm:
      pipeline.run()

    expected_error_prefix = ('Returning a str from a ParDo or FlatMap '
                             'is discouraged.')
    self.assertStartswith(cm.exception.args[0], expected_error_prefix)

  def test_do_with_do_fn_returning_dict_raises_warning(self):
    pipeline = TestPipeline()
    pipeline._options.view_as(TypeOptions).runtime_type_check = True
    pcoll = pipeline | 'Start' >> beam.Create(['2', '9', '3'])
    pcoll | 'Do' >> beam.FlatMap(lambda x: {x: '1'})

    # Since the DoFn directly returns a dict we should get an error warning
    # us.
    with self.assertRaises(typehints.TypeCheckError) as cm:
      pipeline.run()

    expected_error_prefix = ('Returning a dict from a ParDo or FlatMap '
                             'is discouraged.')
    self.assertStartswith(cm.exception.args[0], expected_error_prefix)

  def test_do_with_multiple_outputs_maintains_unique_name(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    r1 = pcoll | 'A' >> beam.FlatMap(lambda x: [x + 1]).with_outputs(main='m')
    r2 = pcoll | 'B' >> beam.FlatMap(lambda x: [x + 2]).with_outputs(main='m')
    assert_that(r1.m, equal_to([2, 3, 4]), label='r1')
    assert_that(r2.m, equal_to([3, 4, 5]), label='r2')
    pipeline.run()

  @attr('ValidatesRunner')
  def test_impulse(self):
    pipeline = TestPipeline()
    result = pipeline | beam.Impulse() | beam.Map(lambda _: 0)
    assert_that(result, equal_to([0]))
    pipeline.run()

  # TODO(BEAM-3544): Disable this test in streaming temporarily.
  # Remove sickbay-streaming tag after it's resolved.
  @attr('ValidatesRunner', 'sickbay-streaming')
  def test_read_metrics(self):
    from apache_beam.io.utils import CountingSource

    class CounterDoFn(beam.DoFn):
      def __init__(self):
        # This counter is unused.
        self.received_records = Metrics.counter(self.__class__,
                                                'receivedRecords')

      def process(self, element):
        self.received_records.inc()

    pipeline = TestPipeline()
    (pipeline | Read(CountingSource(100)) | beam.ParDo(CounterDoFn()))
    res = pipeline.run()
    res.wait_until_finish()
    # This counter is defined in utils.CountingSource.
    metric_results = res.metrics().query(MetricsFilter()
                                         .with_name('recordsRead'))
    outputs_counter = metric_results['counters'][0]
    self.assertEqual(outputs_counter.key.step, 'Read')
    self.assertEqual(outputs_counter.key.metric.name, 'recordsRead')
    self.assertEqual(outputs_counter.committed, 100)
    self.assertEqual(outputs_counter.attempted, 100)

  @attr('ValidatesRunner')
  def test_par_do_with_multiple_outputs_and_using_yield(self):
    class SomeDoFn(beam.DoFn):
      """A custom DoFn using yield."""

      def process(self, element):
        yield element
        if element % 2 == 0:
          yield pvalue.TaggedOutput('even', element)
        else:
          yield pvalue.TaggedOutput('odd', element)

    pipeline = TestPipeline()
    nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
    results = nums | 'ClassifyNumbers' >> beam.ParDo(
        SomeDoFn()).with_outputs('odd', 'even', main='main')
    assert_that(results.main, equal_to([1, 2, 3, 4]))
    assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
    assert_that(results.even, equal_to([2, 4]), label='assert:even')
    pipeline.run()

  @attr('ValidatesRunner')
  def test_par_do_with_multiple_outputs_and_using_return(self):
    def some_fn(v):
      if v % 2 == 0:
        return [v, pvalue.TaggedOutput('even', v)]
      return [v, pvalue.TaggedOutput('odd', v)]

    pipeline = TestPipeline()
    nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
    results = nums | 'ClassifyNumbers' >> beam.FlatMap(
        some_fn).with_outputs('odd', 'even', main='main')
    assert_that(results.main, equal_to([1, 2, 3, 4]))
    assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
    assert_that(results.even, equal_to([2, 4]), label='assert:even')
    pipeline.run()

  @attr('ValidatesRunner')
  def test_undeclared_outputs(self):
    pipeline = TestPipeline()
    nums = pipeline | 'Some Numbers' >> beam.Create([1, 2, 3, 4])
    results = nums | 'ClassifyNumbers' >> beam.FlatMap(
        lambda x: [x,
                   pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x),
                   pvalue.TaggedOutput('extra', x)]
    ).with_outputs()
    assert_that(results[None], equal_to([1, 2, 3, 4]))
    assert_that(results.odd, equal_to([1, 3]), label='assert:odd')
    assert_that(results.even, equal_to([2, 4]), label='assert:even')
    pipeline.run()

  @attr('ValidatesRunner')
  def test_multiple_empty_outputs(self):
    pipeline = TestPipeline()
    nums = pipeline | 'Some Numbers' >> beam.Create([1, 3, 5])
    results = nums | 'ClassifyNumbers' >> beam.FlatMap(
        lambda x: [x,
                   pvalue.TaggedOutput('even' if x % 2 == 0 else 'odd', x)]
    ).with_outputs()
    assert_that(results[None], equal_to([1, 3, 5]))
    assert_that(results.odd, equal_to([1, 3, 5]), label='assert:odd')
    assert_that(results.even, equal_to([]), label='assert:even')
    pipeline.run()

  def test_do_requires_do_fn_returning_iterable(self):
    # This function is incorrect because it returns an object that isn't an
    # iterable.
    def incorrect_par_do_fn(x):
      return x + 5
    pipeline = TestPipeline()
    pipeline._options.view_as(TypeOptions).runtime_type_check = True
    pcoll = pipeline | 'Start' >> beam.Create([2, 9, 3])
    pcoll | 'Do' >> beam.FlatMap(incorrect_par_do_fn)
    # It's a requirement that all user-defined functions to a ParDo return
    # an iterable.
    with self.assertRaises(typehints.TypeCheckError) as cm:
      pipeline.run()

    expected_error_prefix = 'FlatMap and ParDo must return an iterable.'
    self.assertStartswith(cm.exception.args[0], expected_error_prefix)

  def test_do_fn_with_finish(self):
    class MyDoFn(beam.DoFn):
      def process(self, element):
        pass

      def finish_bundle(self):
        yield WindowedValue('finish', -1, [window.GlobalWindow()])

    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    result = pcoll | 'Do' >> beam.ParDo(MyDoFn())

    # May have many bundles, but each has a start and finish.
    def  matcher():
      def match(actual):
        equal_to(['finish'])(list(set(actual)))
        equal_to([1])([actual.count('finish')])
      return match

    assert_that(result, matcher())
    pipeline.run()

  def test_do_fn_with_start(self):
    class MyDoFn(beam.DoFn):
      def __init__(self):
        self.state = 'init'

      def start_bundle(self):
        self.state = 'started'

      def process(self, element):
        if self.state == 'started':
          yield 'started'
        self.state = 'process'

    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3])
    result = pcoll | 'Do' >> beam.ParDo(MyDoFn())

    # May have many bundles, but each has a start and finish.
    def  matcher():
      def match(actual):
        equal_to(['started'])(list(set(actual)))
        equal_to([1])([actual.count('started')])
      return match

    assert_that(result, matcher())
    pipeline.run()

  def test_do_fn_with_start_error(self):
    class MyDoFn(beam.DoFn):
      def start_bundle(self):
        return [1]

      def process(self, element):
        pass

    pipeline = TestPipeline()
    pipeline | 'Start' >> beam.Create([1, 2, 3]) | 'Do' >> beam.ParDo(MyDoFn())
    with self.assertRaises(RuntimeError):
      pipeline.run()

  def test_filter(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([1, 2, 3, 4])
    result = pcoll | 'Filter' >> beam.Filter(lambda x: x % 2 == 0)
    assert_that(result, equal_to([2, 4]))
    pipeline.run()

  class _MeanCombineFn(beam.CombineFn):

    def create_accumulator(self):
      return (0, 0)

    def add_input(self, sum_count, element):
      (sum_, count) = sum_count
      return sum_ + element, count + 1

    def merge_accumulators(self, accumulators):
      sums, counts = zip(*accumulators)
      return sum(sums), sum(counts)

    def extract_output(self, sum_count):
      (sum_, count) = sum_count
      if not count:
        return float('nan')
      return sum_ / float(count)

  def test_combine_with_combine_fn(self):
    vals = [1, 2, 3, 4, 5, 6, 7]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(vals)
    result = pcoll | 'Mean' >> beam.CombineGlobally(self._MeanCombineFn())
    assert_that(result, equal_to([sum(vals) // len(vals)]))
    pipeline.run()

  def test_combine_with_callable(self):
    vals = [1, 2, 3, 4, 5, 6, 7]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(vals)
    result = pcoll | beam.CombineGlobally(sum)
    assert_that(result, equal_to([sum(vals)]))
    pipeline.run()

  def test_combine_with_side_input_as_arg(self):
    values = [1, 2, 3, 4, 5, 6, 7]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(values)
    divisor = pipeline | 'Divisor' >> beam.Create([2])
    result = pcoll | 'Max' >> beam.CombineGlobally(
        # Multiples of divisor only.
        lambda vals, d: max(v for v in vals if v % d == 0),
        pvalue.AsSingleton(divisor)).without_defaults()
    filt_vals = [v for v in values if v % 2 == 0]
    assert_that(result, equal_to([max(filt_vals)]))
    pipeline.run()

  def test_combine_per_key_with_combine_fn(self):
    vals_1 = [1, 2, 3, 4, 5, 6, 7]
    vals_2 = [2, 4, 6, 8, 10, 12, 14]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
                                               [('b', x) for x in vals_2]))
    result = pcoll | 'Mean' >> beam.CombinePerKey(self._MeanCombineFn())
    assert_that(result, equal_to([('a', sum(vals_1) // len(vals_1)),
                                  ('b', sum(vals_2) // len(vals_2))]))
    pipeline.run()

  def test_combine_per_key_with_callable(self):
    vals_1 = [1, 2, 3, 4, 5, 6, 7]
    vals_2 = [2, 4, 6, 8, 10, 12, 14]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
                                               [('b', x) for x in vals_2]))
    result = pcoll | beam.CombinePerKey(sum)
    assert_that(result, equal_to([('a', sum(vals_1)), ('b', sum(vals_2))]))
    pipeline.run()

  def test_combine_per_key_with_side_input_as_arg(self):
    vals_1 = [1, 2, 3, 4, 5, 6, 7]
    vals_2 = [2, 4, 6, 8, 10, 12, 14]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(([('a', x) for x in vals_1] +
                                               [('b', x) for x in vals_2]))
    divisor = pipeline | 'Divisor' >> beam.Create([2])
    result = pcoll | beam.CombinePerKey(
        lambda vals, d: max(v for v in vals if v % d == 0),
        pvalue.AsSingleton(divisor))  # Multiples of divisor only.
    m_1 = max(v for v in vals_1 if v % 2 == 0)
    m_2 = max(v for v in vals_2 if v % 2 == 0)
    assert_that(result, equal_to([('a', m_1), ('b', m_2)]))
    pipeline.run()

  def test_group_by_key(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'start' >> beam.Create(
        [(1, 1), (2, 1), (3, 1), (1, 2), (2, 2), (1, 3)])
    result = pcoll | 'Group' >> beam.GroupByKey()
    assert_that(result, equal_to([(1, [1, 2, 3]), (2, [1, 2]), (3, [1])]))
    pipeline.run()

  def test_partition_with_partition_fn(self):

    class SomePartitionFn(beam.PartitionFn):

      def partition_for(self, element, num_partitions, offset):
        return (element % 3) + offset

    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
    # Attempt nominal partition operation.
    partitions = pcoll | 'Part 1' >> beam.Partition(SomePartitionFn(), 4, 1)
    assert_that(partitions[0], equal_to([]))
    assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')
    assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')
    assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')
    pipeline.run()

    # Check that a bad partition label will yield an error. For the
    # DirectRunner, this error manifests as an exception.
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
    partitions = pcoll | 'Part 2' >> beam.Partition(SomePartitionFn(), 4, 10000)
    with self.assertRaises(ValueError):
      pipeline.run()

  def test_partition_with_callable(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create([0, 1, 2, 3, 4, 5, 6, 7, 8])
    partitions = (
        pcoll | 'part' >> beam.Partition(
            lambda e, n, offset: (e % 3) + offset, 4,
            1))
    assert_that(partitions[0], equal_to([]))
    assert_that(partitions[1], equal_to([0, 3, 6]), label='p1')
    assert_that(partitions[2], equal_to([1, 4, 7]), label='p2')
    assert_that(partitions[3], equal_to([2, 5, 8]), label='p3')
    pipeline.run()

  def test_partition_followed_by_flatten_and_groupbykey(self):
    """Regression test for an issue with how partitions are handled."""
    pipeline = TestPipeline()
    contents = [('aa', 1), ('bb', 2), ('aa', 2)]
    created = pipeline | 'A' >> beam.Create(contents)
    partitioned = created | 'B' >> beam.Partition(lambda x, n: len(x) % n, 3)
    flattened = partitioned | 'C' >> beam.Flatten()
    grouped = flattened | 'D' >> beam.GroupByKey()
    assert_that(grouped, equal_to([('aa', [1, 2]), ('bb', [2])]))
    pipeline.run()

  def test_flatten_pcollections(self):
    pipeline = TestPipeline()
    pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
    pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
    result = (pcoll_1, pcoll_2) | 'Flatten' >> beam.Flatten()
    assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))
    pipeline.run()

  def test_flatten_no_pcollections(self):
    pipeline = TestPipeline()
    with self.assertRaises(ValueError):
      () | 'PipelineArgMissing' >> beam.Flatten()
    result = () | 'Empty' >> beam.Flatten(pipeline=pipeline)
    assert_that(result, equal_to([]))
    pipeline.run()

  def test_flatten_same_pcollections(self):
    pipeline = TestPipeline()
    pc = pipeline | beam.Create(['a', 'b'])
    assert_that((pc, pc, pc) | beam.Flatten(), equal_to(['a', 'b'] * 3))
    pipeline.run()

  def test_flatten_pcollections_in_iterable(self):
    pipeline = TestPipeline()
    pcoll_1 = pipeline | 'Start 1' >> beam.Create([0, 1, 2, 3])
    pcoll_2 = pipeline | 'Start 2' >> beam.Create([4, 5, 6, 7])
    result = [pcoll for pcoll in (pcoll_1, pcoll_2)] | beam.Flatten()
    assert_that(result, equal_to([0, 1, 2, 3, 4, 5, 6, 7]))
    pipeline.run()

  def test_flatten_input_type_must_be_iterable(self):
    # Inputs to flatten *must* be an iterable.
    with self.assertRaises(ValueError):
      4 | beam.Flatten()

  def test_flatten_input_type_must_be_iterable_of_pcolls(self):
    # Inputs to flatten *must* be an iterable of PCollections.
    with self.assertRaises(TypeError):
      {'l': 'test'} | beam.Flatten()
    with self.assertRaises(TypeError):
      set([1, 2, 3]) | beam.Flatten()

  @attr('ValidatesRunner')
  def test_flatten_multiple_pcollections_having_multiple_consumers(self):
    pipeline = TestPipeline()
    input = pipeline | 'Start' >> beam.Create(['AA', 'BBB', 'CC'])

    def split_even_odd(element):
      tag = 'even_length' if len(element) % 2 == 0 else 'odd_length'
      return pvalue.TaggedOutput(tag, element)

    even_length, odd_length = (input | beam.Map(split_even_odd)
                               .with_outputs('even_length', 'odd_length'))
    merged = (even_length, odd_length) | 'Flatten' >> beam.Flatten()

    assert_that(merged, equal_to(['AA', 'BBB', 'CC']))
    assert_that(even_length, equal_to(['AA', 'CC']), label='assert:even')
    assert_that(odd_length, equal_to(['BBB']), label='assert:odd')
    pipeline.run()

  def test_co_group_by_key_on_list(self):
    pipeline = TestPipeline()
    pcoll_1 = pipeline | 'Start 1' >> beam.Create(
        [('a', 1), ('a', 2), ('b', 3), ('c', 4)])
    pcoll_2 = pipeline | 'Start 2' >> beam.Create(
        [('a', 5), ('a', 6), ('c', 7), ('c', 8)])
    result = (pcoll_1, pcoll_2) | beam.CoGroupByKey()
    assert_that(result, equal_to([('a', ([1, 2], [5, 6])),
                                  ('b', ([3], [])),
                                  ('c', ([4], [7, 8]))]))
    pipeline.run()

  def test_co_group_by_key_on_iterable(self):
    pipeline = TestPipeline()
    pcoll_1 = pipeline | 'Start 1' >> beam.Create(
        [('a', 1), ('a', 2), ('b', 3), ('c', 4)])
    pcoll_2 = pipeline | 'Start 2' >> beam.Create(
        [('a', 5), ('a', 6), ('c', 7), ('c', 8)])
    result = [pc for pc in (pcoll_1, pcoll_2)] | beam.CoGroupByKey()
    assert_that(result, equal_to([('a', ([1, 2], [5, 6])),
                                  ('b', ([3], [])),
                                  ('c', ([4], [7, 8]))]))
    pipeline.run()

  def test_co_group_by_key_on_dict(self):
    pipeline = TestPipeline()
    pcoll_1 = pipeline | 'Start 1' >> beam.Create(
        [('a', 1), ('a', 2), ('b', 3), ('c', 4)])
    pcoll_2 = pipeline | 'Start 2' >> beam.Create(
        [('a', 5), ('a', 6), ('c', 7), ('c', 8)])
    result = {'X': pcoll_1, 'Y': pcoll_2} | beam.CoGroupByKey()
    assert_that(result, equal_to([('a', {'X': [1, 2], 'Y': [5, 6]}),
                                  ('b', {'X': [3], 'Y': []}),
                                  ('c', {'X': [4], 'Y': [7, 8]})]))
    pipeline.run()

  def test_group_by_key_input_must_be_kv_pairs(self):
    pipeline = TestPipeline()
    pcolls = pipeline | 'A' >> beam.Create([1, 2, 3, 4, 5])

    with self.assertRaises(typehints.TypeCheckError) as e:
      pcolls | 'D' >> beam.GroupByKey()
      pipeline.run()

    self.assertStartswith(
        e.exception.args[0],
        'Input type hint violation at D: expected '
        'Tuple[TypeVariable[K], TypeVariable[V]]')

  def test_group_by_key_only_input_must_be_kv_pairs(self):
    pipeline = TestPipeline()
    pcolls = pipeline | 'A' >> beam.Create(['a', 'b', 'f'])
    with self.assertRaises(typehints.TypeCheckError) as cm:
      pcolls | 'D' >> _GroupByKeyOnly()
      pipeline.run()

    expected_error_prefix = ('Input type hint violation at D: expected '
                             'Tuple[TypeVariable[K], TypeVariable[V]]')
    self.assertStartswith(cm.exception.args[0], expected_error_prefix)

  def test_keys_and_values(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(
        [(3, 1), (2, 1), (1, 1), (3, 2), (2, 2), (3, 3)])
    keys = pcoll.apply(beam.Keys('keys'))
    vals = pcoll.apply(beam.Values('vals'))
    assert_that(keys, equal_to([1, 2, 2, 3, 3, 3]), label='assert:keys')
    assert_that(vals, equal_to([1, 1, 1, 2, 2, 3]), label='assert:vals')
    pipeline.run()

  def test_kv_swap(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(
        [(6, 3), (1, 2), (7, 1), (5, 2), (3, 2)])
    result = pcoll.apply(beam.KvSwap(), label='swap')
    assert_that(result, equal_to([(1, 7), (2, 1), (2, 3), (2, 5), (3, 6)]))
    pipeline.run()

  def test_distinct(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(
        [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])
    result = pcoll.apply(beam.Distinct())
    assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))
    pipeline.run()

  def test_remove_duplicates(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(
        [6, 3, 1, 1, 9, 'pleat', 'pleat', 'kazoo', 'navel'])
    result = pcoll.apply(beam.RemoveDuplicates())
    assert_that(result, equal_to([1, 3, 6, 9, 'pleat', 'kazoo', 'navel']))
    pipeline.run()

  def test_chained_ptransforms(self):
    pipeline = TestPipeline()
    t = (beam.Map(lambda x: (x, 1))
         | beam.GroupByKey()
         | beam.Map(lambda x_ones: (x_ones[0], sum(x_ones[1]))))
    result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t
    assert_that(result, equal_to([('a', 2), ('b', 1)]))
    pipeline.run()

  def test_apply_to_list(self):
    self.assertCountEqual(
        [1, 2, 3], [0, 1, 2] | 'AddOne' >> beam.Map(lambda x: x + 1))
    self.assertCountEqual([1],
                          [0, 1, 2] | 'Odd' >> beam.Filter(lambda x: x % 2))
    self.assertCountEqual([1, 2, 100, 3],
                          ([1, 2, 3], [100]) | beam.Flatten())
    join_input = ([('k', 'a')],
                  [('k', 'b'), ('k', 'c')])
    self.assertCountEqual([('k', (['a'], ['b', 'c']))],
                          join_input | beam.CoGroupByKey())

  def test_multi_input_ptransform(self):
    class DisjointUnion(PTransform):
      def expand(self, pcollections):
        return (pcollections
                | beam.Flatten()
                | beam.Map(lambda x: (x, None))
                | beam.GroupByKey()
                | beam.Map(lambda kv: kv[0]))
    self.assertEqual([1, 2, 3], sorted(([1, 2], [2, 3]) | DisjointUnion()))

  def test_apply_to_crazy_pvaluish(self):
    class NestedFlatten(PTransform):
      """A PTransform taking and returning nested PValueish.

      Takes as input a list of dicts, and returns a dict with the corresponding
      values flattened.
      """
      def _extract_input_pvalues(self, pvalueish):
        pvalueish = list(pvalueish)
        return pvalueish, sum([list(p.values()) for p in pvalueish], [])

      def expand(self, pcoll_dicts):
        keys = reduce(operator.or_, [set(p.keys()) for p in pcoll_dicts])
        res = {}
        for k in keys:
          res[k] = [p[k] for p in pcoll_dicts if k in p] | k >> beam.Flatten()
        return res
    res = [{'a': [1, 2, 3]},
           {'a': [4, 5, 6], 'b': ['x', 'y', 'z']},
           {'a': [7, 8], 'b': ['x', 'y'], 'c': []}] | NestedFlatten()
    self.assertEqual(3, len(res))
    self.assertEqual([1, 2, 3, 4, 5, 6, 7, 8], sorted(res['a']))
    self.assertEqual(['x', 'x', 'y', 'y', 'z'], sorted(res['b']))
    self.assertEqual([], sorted(res['c']))

  def test_named_tuple(self):
    MinMax = collections.namedtuple('MinMax', ['min', 'max'])

    class MinMaxTransform(PTransform):
      def expand(self, pcoll):
        return MinMax(
            min=pcoll | beam.CombineGlobally(min).without_defaults(),
            max=pcoll | beam.CombineGlobally(max).without_defaults())
    res = [1, 2, 4, 8] | MinMaxTransform()
    self.assertIsInstance(res, MinMax)
    self.assertEqual(res, MinMax(min=[1], max=[8]))

    flat = res | beam.Flatten()
    self.assertEqual(sorted(flat), [1, 8])

  def test_tuple_twice(self):
    class Duplicate(PTransform):
      def expand(self, pcoll):
        return pcoll, pcoll

    res1, res2 = [1, 2, 4, 8] | Duplicate()
    self.assertEqual(sorted(res1), [1, 2, 4, 8])
    self.assertEqual(sorted(res2), [1, 2, 4, 8])


@beam.ptransform_fn
def SamplePTransform(pcoll):
  """Sample transform using the @ptransform_fn decorator."""
  map_transform = 'ToPairs' >> beam.Map(lambda v: (v, None))
  combine_transform = 'Group' >> beam.CombinePerKey(lambda vs: None)
  keys_transform = 'Distinct' >> beam.Keys()
  return pcoll | map_transform | combine_transform | keys_transform


class PTransformLabelsTest(unittest.TestCase):

  class CustomTransform(beam.PTransform):

    pardo = None

    def expand(self, pcoll):
      self.pardo = '*Do*' >> beam.FlatMap(lambda x: [x + 1])
      return pcoll | self.pardo

  def test_chained_ptransforms(self):
    """Tests that chaining gets proper nesting."""
    pipeline = TestPipeline()
    map1 = 'Map1' >> beam.Map(lambda x: (x, 1))
    gbk = 'Gbk' >> beam.GroupByKey()
    map2 = 'Map2' >> beam.Map(lambda x_ones2: (x_ones2[0], sum(x_ones2[1])))
    t = (map1 | gbk | map2)
    result = pipeline | 'Start' >> beam.Create(['a', 'a', 'b']) | t
    self.assertTrue('Map1|Gbk|Map2/Map1' in pipeline.applied_labels)
    self.assertTrue('Map1|Gbk|Map2/Gbk' in pipeline.applied_labels)
    self.assertTrue('Map1|Gbk|Map2/Map2' in pipeline.applied_labels)
    assert_that(result, equal_to([('a', 2), ('b', 1)]))
    pipeline.run()

  def test_apply_custom_transform_without_label(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
    custom = PTransformLabelsTest.CustomTransform()
    result = pipeline.apply(custom, pcoll)
    self.assertTrue('CustomTransform' in pipeline.applied_labels)
    self.assertTrue('CustomTransform/*Do*' in pipeline.applied_labels)
    assert_that(result, equal_to([2, 3, 4]))
    pipeline.run()

  def test_apply_custom_transform_with_label(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
    custom = PTransformLabelsTest.CustomTransform('*Custom*')
    result = pipeline.apply(custom, pcoll)
    self.assertTrue('*Custom*' in pipeline.applied_labels)
    self.assertTrue('*Custom*/*Do*' in pipeline.applied_labels)
    assert_that(result, equal_to([2, 3, 4]))
    pipeline.run()

  def test_combine_without_label(self):
    vals = [1, 2, 3, 4, 5, 6, 7]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(vals)
    combine = beam.CombineGlobally(sum)
    result = pcoll | combine
    self.assertTrue('CombineGlobally(sum)' in pipeline.applied_labels)
    assert_that(result, equal_to([sum(vals)]))
    pipeline.run()

  def test_apply_ptransform_using_decorator(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'PColl' >> beam.Create([1, 2, 3])
    _ = pcoll | '*Sample*' >> SamplePTransform()
    self.assertTrue('*Sample*' in pipeline.applied_labels)
    self.assertTrue('*Sample*/ToPairs' in pipeline.applied_labels)
    self.assertTrue('*Sample*/Group' in pipeline.applied_labels)
    self.assertTrue('*Sample*/Distinct' in pipeline.applied_labels)

  def test_combine_with_label(self):
    vals = [1, 2, 3, 4, 5, 6, 7]
    pipeline = TestPipeline()
    pcoll = pipeline | 'Start' >> beam.Create(vals)
    combine = '*Sum*' >> beam.CombineGlobally(sum)
    result = pcoll | combine
    self.assertTrue('*Sum*' in pipeline.applied_labels)
    assert_that(result, equal_to([sum(vals)]))
    pipeline.run()

  def check_label(self, ptransform, expected_label):
    pipeline = TestPipeline()
    pipeline | 'Start' >> beam.Create([('a', 1)]) | ptransform
    actual_label = sorted(pipeline.applied_labels - {'Start', 'Start/Read'})[0]
    self.assertEqual(expected_label, re.sub(r'\d{3,}', '#', actual_label))

  def test_default_labels(self):
    self.check_label(beam.Map(len), r'Map(len)')
    self.check_label(beam.Map(lambda x: x),
                     r'Map(<lambda at ptransform_test.py:#>)')
    self.check_label(beam.FlatMap(list), r'FlatMap(list)')
    self.check_label(beam.Filter(sum), r'Filter(sum)')
    self.check_label(beam.CombineGlobally(sum), r'CombineGlobally(sum)')
    self.check_label(beam.CombinePerKey(sum), r'CombinePerKey(sum)')

    class MyDoFn(beam.DoFn):
      def process(self, unused_element):
        pass

    self.check_label(beam.ParDo(MyDoFn()), r'ParDo(MyDoFn)')

  def test_label_propogation(self):
    self.check_label('TestMap' >> beam.Map(len), r'TestMap')
    self.check_label('TestLambda' >> beam.Map(lambda x: x), r'TestLambda')
    self.check_label('TestFlatMap' >> beam.FlatMap(list), r'TestFlatMap')
    self.check_label('TestFilter' >> beam.Filter(sum), r'TestFilter')
    self.check_label('TestCG' >> beam.CombineGlobally(sum), r'TestCG')
    self.check_label('TestCPK' >> beam.CombinePerKey(sum), r'TestCPK')

    class MyDoFn(beam.DoFn):
      def process(self, unused_element):
        pass

    self.check_label('TestParDo' >> beam.ParDo(MyDoFn()), r'TestParDo')


class PTransformTestDisplayData(unittest.TestCase):
  def test_map_named_function(self):
    tr = beam.Map(len)
    dd = DisplayData.create_from(tr)
    nspace = 'apache_beam.transforms.core.CallableWrapperDoFn'
    expected_item = DisplayDataItem('len', key='fn',
                                    label='Transform Function',
                                    namespace=nspace)
    hc.assert_that(dd.items, hc.has_item(expected_item))

  def test_map_anonymous_function(self):
    tr = beam.Map(lambda x: x)
    dd = DisplayData.create_from(tr)
    nspace = 'apache_beam.transforms.core.CallableWrapperDoFn'
    expected_item = DisplayDataItem('<lambda>', key='fn',
                                    label='Transform Function',
                                    namespace=nspace)
    hc.assert_that(dd.items, hc.has_item(expected_item))

  def test_flatmap_named_function(self):
    tr = beam.FlatMap(list)
    dd = DisplayData.create_from(tr)
    nspace = 'apache_beam.transforms.core.CallableWrapperDoFn'
    expected_item = DisplayDataItem('list', key='fn',
                                    label='Transform Function',
                                    namespace=nspace)
    hc.assert_that(dd.items, hc.has_item(expected_item))

  def test_flatmap_anonymous_function(self):
    tr = beam.FlatMap(lambda x: [x])
    dd = DisplayData.create_from(tr)
    nspace = 'apache_beam.transforms.core.CallableWrapperDoFn'
    expected_item = DisplayDataItem('<lambda>', key='fn',
                                    label='Transform Function',
                                    namespace=nspace)
    hc.assert_that(dd.items, hc.has_item(expected_item))

  def test_filter_named_function(self):
    tr = beam.Filter(sum)
    dd = DisplayData.create_from(tr)
    nspace = 'apache_beam.transforms.core.CallableWrapperDoFn'
    expected_item = DisplayDataItem('sum', key='fn',
                                    label='Transform Function',
                                    namespace=nspace)
    hc.assert_that(dd.items, hc.has_item(expected_item))

  def test_filter_anonymous_function(self):
    tr = beam.Filter(lambda x: x // 30)
    dd = DisplayData.create_from(tr)
    nspace = 'apache_beam.transforms.core.CallableWrapperDoFn'
    expected_item = DisplayDataItem('<lambda>', key='fn',
                                    label='Transform Function',
                                    namespace=nspace)
    hc.assert_that(dd.items, hc.has_item(expected_item))


class PTransformTypeCheckTestCase(TypeHintTestCase):

  def assertStartswith(self, msg, prefix):
    self.assertTrue(msg.startswith(prefix),
                    '"%s" does not start with "%s"' % (msg, prefix))

  def setUp(self):
    self.p = TestPipeline()

  def test_do_fn_pipeline_pipeline_type_check_satisfied(self):
    @with_input_types(int, int)
    @with_output_types(int)
    class AddWithFive(beam.DoFn):
      def process(self, element, five):
        return [element + five]

    d = (self.p
         | 'T' >> beam.Create([1, 2, 3]).with_output_types(int)
         | 'Add' >> beam.ParDo(AddWithFive(), 5))

    assert_that(d, equal_to([6, 7, 8]))
    self.p.run()

  def test_do_fn_pipeline_pipeline_type_check_violated(self):
    @with_input_types(str, str)
    @with_output_types(str)
    class ToUpperCaseWithPrefix(beam.DoFn):
      def process(self, element, prefix):
        return [prefix + element.upper()]

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'T' >> beam.Create([1, 2, 3]).with_output_types(int)
       | 'Upper' >> beam.ParDo(ToUpperCaseWithPrefix(), 'hello'))

    self.assertEqual("Type hint violation for 'Upper': "
                     "requires {} but got {} for element".format(str, int),
                     e.exception.args[0])

  def test_do_fn_pipeline_runtime_type_check_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    @with_input_types(int, int)
    @with_output_types(int)
    class AddWithNum(beam.DoFn):
      def process(self, element, num):
        return [element + num]

    d = (self.p
         | 'T' >> beam.Create([1, 2, 3]).with_output_types(int)
         | 'Add' >> beam.ParDo(AddWithNum(), 5))

    assert_that(d, equal_to([6, 7, 8]))
    self.p.run()

  def test_do_fn_pipeline_runtime_type_check_violated(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    @with_input_types(int, int)
    @with_output_types(int)
    class AddWithNum(beam.DoFn):
      def process(self, element, num):
        return [element + num]

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'T' >> beam.Create(['1', '2', '3']).with_output_types(str)
       | 'Add' >> beam.ParDo(AddWithNum(), 5))
      self.p.run()

    self.assertEqual("Type hint violation for 'Add': "
                     "requires {} but got {} for element".format(int, str),
                     e.exception.args[0])

  def test_pardo_does_not_type_check_using_type_hint_decorators(self):
    @with_input_types(a=int)
    @with_output_types(typing.List[str])
    def int_to_str(a):
      return [str(a)]

    # The function above is expecting an int for its only parameter. However, it
    # will receive a str instead, which should result in a raised exception.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'S' >> beam.Create(['b', 'a', 'r']).with_output_types(str)
       | 'ToStr' >> beam.FlatMap(int_to_str))

    self.assertEqual("Type hint violation for 'ToStr': "
                     "requires {} but got {} for a".format(int, str),
                     e.exception.args[0])

  def test_pardo_properly_type_checks_using_type_hint_decorators(self):
    @with_input_types(a=str)
    @with_output_types(typing.List[str])
    def to_all_upper_case(a):
      return [a.upper()]

    # If this type-checks than no error should be raised.
    d = (self.p
         | 'T' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str)
         | 'Case' >> beam.FlatMap(to_all_upper_case))
    assert_that(d, equal_to(['T', 'E', 'S', 'T']))
    self.p.run()

    # Output type should have been recognized as 'str' rather than List[str] to
    # do the flatten part of FlatMap.
    self.assertEqual(str, d.element_type)

  def test_pardo_does_not_type_check_using_type_hint_methods(self):
    # The first ParDo outputs pcoll's of type int, however the second ParDo is
    # expecting pcoll's of type str instead.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'S' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str)
       | ('Score' >> beam.FlatMap(lambda x: [1] if x == 't' else [2])
          .with_input_types(str).with_output_types(int))
       | ('Upper' >> beam.FlatMap(lambda x: [x.upper()])
          .with_input_types(str).with_output_types(str)))

    self.assertEqual("Type hint violation for 'Upper': "
                     "requires {} but got {} for x".format(str, int),
                     e.exception.args[0])

  def test_pardo_properly_type_checks_using_type_hint_methods(self):
    # Pipeline should be created successfully without an error
    d = (self.p
         | 'S' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str)
         | 'Dup' >> beam.FlatMap(lambda x: [x + x])
         .with_input_types(str).with_output_types(str)
         | 'Upper' >> beam.FlatMap(lambda x: [x.upper()])
         .with_input_types(str).with_output_types(str))

    assert_that(d, equal_to(['TT', 'EE', 'SS', 'TT']))
    self.p.run()

  def test_map_does_not_type_check_using_type_hints_methods(self):
    # The transform before 'Map' has indicated that it outputs PCollections with
    # int's, while Map is expecting one of str.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'S' >> beam.Create([1, 2, 3, 4]).with_output_types(int)
       | 'Upper' >> beam.Map(lambda x: x.upper())
       .with_input_types(str).with_output_types(str))

    self.assertEqual("Type hint violation for 'Upper': "
                     "requires {} but got {} for x".format(str, int),
                     e.exception.args[0])

  def test_map_properly_type_checks_using_type_hints_methods(self):
    # No error should be raised if this type-checks properly.
    d = (self.p
         | 'S' >> beam.Create([1, 2, 3, 4]).with_output_types(int)
         | 'ToStr' >> beam.Map(lambda x: str(x))
         .with_input_types(int).with_output_types(str))
    assert_that(d, equal_to(['1', '2', '3', '4']))
    self.p.run()

  def test_map_does_not_type_check_using_type_hints_decorator(self):
    @with_input_types(s=str)
    @with_output_types(str)
    def upper(s):
      return s.upper()

    # Hinted function above expects a str at pipeline construction.
    # However, 'Map' should detect that Create has hinted an int instead.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'S' >> beam.Create([1, 2, 3, 4]).with_output_types(int)
       | 'Upper' >> beam.Map(upper))

    self.assertEqual("Type hint violation for 'Upper': "
                     "requires {} but got {} for s".format(str, int),
                     e.exception.args[0])

  def test_map_properly_type_checks_using_type_hints_decorator(self):
    @with_input_types(a=bool)
    @with_output_types(int)
    def bool_to_int(a):
      return int(a)

    # If this type-checks than no error should be raised.
    d = (self.p
         | 'Bools' >> beam.Create([True, False, True]).with_output_types(bool)
         | 'ToInts' >> beam.Map(bool_to_int))
    assert_that(d, equal_to([1, 0, 1]))
    self.p.run()

  def test_filter_does_not_type_check_using_type_hints_method(self):
    # Filter is expecting an int but instead looks to the 'left' and sees a str
    # incoming.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'Strs' >> beam.Create(['1', '2', '3', '4', '5']).with_output_types(str)
       | 'Lower' >> beam.Map(lambda x: x.lower())
       .with_input_types(str).with_output_types(str)
       | 'Below 3' >> beam.Filter(lambda x: x < 3).with_input_types(int))

    self.assertEqual("Type hint violation for 'Below 3': "
                     "requires {} but got {} for x".format(int, str),
                     e.exception.args[0])

  def test_filter_type_checks_using_type_hints_method(self):
    # No error should be raised if this type-checks properly.
    d = (self.p
         | beam.Create(['1', '2', '3', '4', '5']).with_output_types(str)
         | 'ToInt' >> beam.Map(lambda x: int(x))
         .with_input_types(str).with_output_types(int)
         | 'Below 3' >> beam.Filter(lambda x: x < 3).with_input_types(int))
    assert_that(d, equal_to([1, 2]))
    self.p.run()

  def test_filter_does_not_type_check_using_type_hints_decorator(self):
    @with_input_types(a=float)
    def more_than_half(a):
      return a > 0.50

    # Func above was hinted to only take a float, yet an int will be passed.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'Ints' >> beam.Create([1, 2, 3, 4]).with_output_types(int)
       | 'Half' >> beam.Filter(more_than_half))

    self.assertEqual("Type hint violation for 'Half': "
                     "requires {} but got {} for a".format(float, int),
                     e.exception.args[0])

  def test_filter_type_checks_using_type_hints_decorator(self):
    @with_input_types(b=int)
    def half(b):
      import random
      return bool(random.choice([0, 1]))

    # Filter should deduce that it returns the same type that it takes.
    (self.p
     | 'Str' >> beam.Create(range(5)).with_output_types(int)
     | 'Half' >> beam.Filter(half)
     | 'ToBool' >> beam.Map(lambda x: bool(x))
     .with_input_types(int).with_output_types(bool))

  def test_group_by_key_only_output_type_deduction(self):
    d = (self.p
         | 'Str' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str)
         | ('Pair' >> beam.Map(lambda x: (x, ord(x)))
            .with_output_types(typing.Tuple[str, str]))
         | _GroupByKeyOnly())

    # Output type should correctly be deduced.
    # GBK-only should deduce that Tuple[A, B] is turned into
    # Tuple[A, Iterable[B]].
    self.assertCompatible(typing.Tuple[str, typing.Iterable[str]],
                          d.element_type)

  def test_group_by_key_output_type_deduction(self):
    d = (self.p
         | 'Str' >> beam.Create(range(20)).with_output_types(int)
         | ('PairNegative' >> beam.Map(lambda x: (x % 5, -x))
            .with_output_types(typing.Tuple[int, int]))
         | beam.GroupByKey())

    # Output type should correctly be deduced.
    # GBK should deduce that Tuple[A, B] is turned into Tuple[A, Iterable[B]].
    self.assertCompatible(typing.Tuple[int, typing.Iterable[int]],
                          d.element_type)

  def test_group_by_key_only_does_not_type_check(self):
    # GBK will be passed raw int's here instead of some form of Tuple[A, B].
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([1, 2, 3]).with_output_types(int)
       | 'F' >> _GroupByKeyOnly())

    self.assertEqual("Input type hint violation at F: "
                     "expected Tuple[TypeVariable[K], TypeVariable[V]], "
                     "got {}".format(int),
                     e.exception.args[0])

  def test_group_by_does_not_type_check(self):
    # Create is returning a List[int, str], rather than a Tuple[int, str]
    # that is aliased to Tuple[int, str].
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | (beam.Create([[1], [2]])
          .with_output_types(typing.Iterable[int]))
       | 'T' >> beam.GroupByKey())

    self.assertEqual("Input type hint violation at T: "
                     "expected Tuple[TypeVariable[K], TypeVariable[V]], "
                     "got Iterable[int]",
                     e.exception.args[0])

  def test_pipeline_checking_pardo_insufficient_type_information(self):
    self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED'

    # Type checking is enabled, but 'Create' doesn't pass on any relevant type
    # information to the ParDo.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'Nums' >> beam.Create(range(5))
       | 'ModDup' >> beam.FlatMap(lambda x: (x % 2, x)))

    self.assertEqual('Pipeline type checking is enabled, however no output '
                     'type-hint was found for the PTransform Create(Nums)',
                     e.exception.args[0])

  def test_pipeline_checking_gbk_insufficient_type_information(self):
    self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED'
    # Type checking is enabled, but 'Map' doesn't pass on any relevant type
    # information to GBK-only.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'Nums' >> beam.Create(range(5)).with_output_types(int)
       | 'ModDup' >> beam.Map(lambda x: (x % 2, x))
       | _GroupByKeyOnly())

    self.assertEqual('Pipeline type checking is enabled, however no output '
                     'type-hint was found for the PTransform '
                     'ParDo(ModDup)',
                     e.exception.args[0])

  def test_disable_pipeline_type_check(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    # The pipeline below should raise a TypeError, however pipeline type
    # checking was disabled above.
    (self.p
     | 'T' >> beam.Create([1, 2, 3]).with_output_types(int)
     | 'Lower' >> beam.Map(lambda x: x.lower())
     .with_input_types(str).with_output_types(str))

  def test_run_time_type_checking_enabled_type_violation(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    @with_output_types(str)
    @with_input_types(x=int)
    def int_to_string(x):
      return str(x)

    # Function above has been type-hinted to only accept an int. But in the
    # pipeline execution it'll be passed a string due to the output of Create.
    (self.p
     | 'T' >> beam.Create(['some_string'])
     | 'ToStr' >> beam.Map(int_to_string))
    with self.assertRaises(typehints.TypeCheckError) as e:
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within ParDo(ToStr): "
        "Type-hint for argument: 'x' violated. "
        "Expected an instance of {}, "
        "instead found some_string, an instance of {}.".format(int, str))

  def test_run_time_type_checking_enabled_types_satisfied(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    @with_output_types(typing.Tuple[int, str])
    @with_input_types(x=str)
    def group_with_upper_ord(x):
      return (ord(x.upper()) % 5, x)

    # Pipeline checking is off, but the above function should satisfy types at
    # run-time.
    result = (self.p
              | 'T' >> beam.Create(['t', 'e', 's', 't', 'i', 'n', 'g'])
              .with_output_types(str)
              | 'GenKeys' >> beam.Map(group_with_upper_ord)
              | 'O' >> beam.GroupByKey())

    assert_that(result, equal_to([(1, ['g']),
                                  (3, ['s', 'i', 'n']),
                                  (4, ['t', 'e', 't'])]))
    self.p.run()

  def test_pipeline_checking_satisfied_but_run_time_types_violate(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    @with_output_types(typing.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(5)).with_output_types(int)
     | 'IsEven' >> beam.Map(is_even_as_key)
     | 'Parity' >> beam.GroupByKey())

    # Although all the types appear to be correct when checked at pipeline
    # construction. Runtime type-checking should detect the 'is_even_as_key' is
    # returning Tuple[int, int], instead of Tuple[bool, int].
    with self.assertRaises(typehints.TypeCheckError) as e:
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within ParDo(IsEven): "
        "Tuple[bool, int] hint type-constraint violated. "
        "The type of element #0 in the passed tuple is incorrect. "
        "Expected an instance of type bool, "
        "instead received an instance of type int.")

  def test_pipeline_checking_satisfied_run_time_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    @with_output_types(typing.Tuple[bool, int])
    @with_input_types(a=int)
    def is_even_as_key(a):
      # The programming error in the above test-case has now been fixed.
      # Everything should properly type-check.
      return (a % 2 == 0, a)

    result = (self.p
              | 'Nums' >> beam.Create(range(5)).with_output_types(int)
              | 'IsEven' >> beam.Map(is_even_as_key)
              | 'Parity' >> beam.GroupByKey())

    assert_that(result, equal_to([(False, [1, 3]), (True, [0, 2, 4])]))
    self.p.run()

  def test_pipeline_runtime_checking_violation_simple_type_input(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    # The type-hinted applied via the 'with_input_types()' method indicates the
    # ParDo should receive an instance of type 'str', however an 'int' will be
    # passed instead.
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([1, 2, 3])
       | ('ToInt' >> beam.FlatMap(lambda x: [int(x)])
          .with_input_types(str).with_output_types(int)))
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within ParDo(ToInt): "
        "Type-hint for argument: 'x' violated. "
        "Expected an instance of {}, "
        "instead found 1, an instance of {}.".format(str, int))

  def test_pipeline_runtime_checking_violation_composite_type_input(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([(1, 3.0), (2, 4.9), (3, 9.5)])
       | ('Add' >> beam.FlatMap(lambda x_y: [x_y[0] + x_y[1]])
          .with_input_types(typing.Tuple[int, int]).with_output_types(int))
      )
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within ParDo(Add): "
        "Type-hint for argument: 'x_y' violated: "
        "Tuple[int, int] hint type-constraint violated. "
        "The type of element #1 in the passed tuple is incorrect. "
        "Expected an instance of type int, instead received an instance "
        "of type float.")

  def test_pipeline_runtime_checking_violation_simple_type_output(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    # The type-hinted applied via the 'returns()' method indicates the ParDo
    # should output an instance of type 'int', however a 'float' will be
    # generated instead.
    print("HINTS", ('ToInt' >> beam.FlatMap(
        lambda x: [float(x)]).with_input_types(int).with_output_types(
            int)).get_type_hints())
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([1, 2, 3])
       | ('ToInt' >> beam.FlatMap(lambda x: [float(x)])
          .with_input_types(int).with_output_types(int))
      )
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within "
        "ParDo(ToInt): "
        "According to type-hint expected output should be "
        "of type {}. Instead, received '1.0', "
        "an instance of type {}.".format(int, float))

  def test_pipeline_runtime_checking_violation_composite_type_output(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True
    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.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([(1, 3.0), (2, 4.9), (3, 9.5)])
       | ('Swap' >> beam.FlatMap(lambda x_y1: [x_y1[0] + x_y1[1]])
          .with_input_types(typing.Tuple[int, float])
          .with_output_types(typing.Tuple[float, int]))
      )
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within "
        "ParDo(Swap): Tuple type constraint violated. "
        "Valid object instance must be of type 'tuple'. Instead, "
        "an instance of 'float' was received.")

  def test_pipeline_runtime_checking_violation_with_side_inputs_decorator(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    @with_output_types(int)
    @with_input_types(a=int, b=int)
    def add(a, b):
      return a + b

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p | beam.Create([1, 2, 3, 4]) | 'Add 1' >> beam.Map(add, 1.0))
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within ParDo(Add 1): "
        "Type-hint for argument: 'b' violated. "
        "Expected an instance of {}, "
        "instead found 1.0, an instance of {}.".format(int, float))

  def test_pipeline_runtime_checking_violation_with_side_inputs_via_method(self):  # pylint: disable=line-too-long
    self.p._options.view_as(TypeOptions).runtime_type_check = True
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([1, 2, 3, 4])
       | ('Add 1' >> beam.Map(lambda x, one: x + one, 1.0)
          .with_input_types(int, int)
          .with_output_types(float)))
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within ParDo(Add 1): "
        "Type-hint for argument: 'one' violated. "
        "Expected an instance of {}, "
        "instead found 1.0, an instance of {}.".format(int, float))

  def test_combine_properly_pipeline_type_checks_using_decorator(self):
    @with_output_types(int)
    @with_input_types(ints=typing.Iterable[int])
    def sum_ints(ints):
      return sum(ints)

    d = (self.p
         | 'T' >> beam.Create([1, 2, 3]).with_output_types(int)
         | 'Sum' >> beam.CombineGlobally(sum_ints))

    self.assertEqual(int, d.element_type)
    assert_that(d, equal_to([6]))
    self.p.run()

  def test_combine_func_type_hint_does_not_take_iterable_using_decorator(self):
    @with_output_types(int)
    @with_input_types(a=int)
    def bad_combine(a):
      5 + a

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'M' >> beam.Create([1, 2, 3]).with_output_types(int)
       | 'Add' >> beam.CombineGlobally(bad_combine))

    self.assertEqual(
        "All functions for a Combine PTransform must accept a "
        "single argument compatible with: Iterable[Any]. "
        "Instead a function with input type: {} was received.".format(int),
        e.exception.args[0])

  def test_combine_pipeline_type_propagation_using_decorators(self):
    @with_output_types(int)
    @with_input_types(ints=typing.Iterable[int])
    def sum_ints(ints):
      return sum(ints)

    @with_output_types(typing.List[int])
    @with_input_types(n=int)
    def range_from_zero(n):
      return list(range(n+1))

    d = (self.p
         | 'T' >> beam.Create([1, 2, 3]).with_output_types(int)
         | 'Sum' >> beam.CombineGlobally(sum_ints)
         | 'Range' >> beam.ParDo(range_from_zero))

    self.assertEqual(int, d.element_type)
    assert_that(d, equal_to([0, 1, 2, 3, 4, 5, 6]))
    self.p.run()

  def test_combine_runtime_type_check_satisfied_using_decorators(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False

    @with_output_types(int)
    @with_input_types(ints=typing.Iterable[int])
    def iter_mul(ints):
      return reduce(operator.mul, ints, 1)

    d = (self.p
         | 'K' >> beam.Create([5, 5, 5, 5]).with_output_types(int)
         | 'Mul' >> beam.CombineGlobally(iter_mul))

    assert_that(d, equal_to([625]))
    self.p.run()

  def test_combine_runtime_type_check_violation_using_decorators(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    # Combine fn is returning the incorrect type
    @with_output_types(int)
    @with_input_types(ints=typing.Iterable[int])
    def iter_mul(ints):
      return str(reduce(operator.mul, ints, 1))

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'K' >> beam.Create([5, 5, 5, 5]).with_output_types(int)
       | 'Mul' >> beam.CombineGlobally(iter_mul))
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within "
        "Mul/CombinePerKey: "
        "Type-hint for return type violated. "
        "Expected an instance of {}, instead found".format(int))

  def test_combine_pipeline_type_check_using_methods(self):
    d = (self.p
         | beam.Create(['t', 'e', 's', 't']).with_output_types(str)
         | ('concat' >> beam.CombineGlobally(lambda s: ''.join(s))
            .with_input_types(str).with_output_types(str)))

    def matcher(expected):
      def match(actual):
        equal_to(expected)(list(actual[0]))
      return match
    assert_that(d, matcher('estt'))
    self.p.run()

  def test_combine_runtime_type_check_using_methods(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create(range(5)).with_output_types(int)
         | ('Sum' >> beam.CombineGlobally(lambda s: sum(s))
            .with_input_types(int).with_output_types(int)))

    assert_that(d, equal_to([10]))
    self.p.run()

  def test_combine_pipeline_type_check_violation_using_methods(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create(range(3)).with_output_types(int)
       | ('SortJoin' >> beam.CombineGlobally(lambda s: ''.join(sorted(s)))
          .with_input_types(str).with_output_types(str)))

    self.assertEqual("Input type hint violation at SortJoin: "
                     "expected {}, got {}".format(str, int),
                     e.exception.args[0])

  def test_combine_runtime_type_check_violation_using_methods(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create(range(3)).with_output_types(int)
       | ('SortJoin' >> beam.CombineGlobally(lambda s: ''.join(sorted(s)))
          .with_input_types(str).with_output_types(str)))
      self.p.run()

    self.assertStartswith(
        e.exception.args[0],
        "Runtime type violation detected within "
        "ParDo(SortJoin/KeyWithVoid): "
        "Type-hint for argument: 'v' violated. "
        "Expected an instance of {}, "
        "instead found 0, an instance of {}.".format(str, int))

  def test_combine_insufficient_type_hint_information(self):
    self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED'

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'E' >> beam.Create(range(3)).with_output_types(int)
       | 'SortJoin' >> beam.CombineGlobally(lambda s: ''.join(sorted(s)))
       | 'F' >> beam.Map(lambda x: x + 1))

    self.assertStartswith(
        e.exception.args[0],
        'Pipeline type checking is enabled, '
        'however no output type-hint was found for the PTransform '
        'ParDo('
        'SortJoin/CombinePerKey/')

  def test_mean_globally_pipeline_checking_satisfied(self):
    d = (self.p
         | 'C' >> beam.Create(range(5)).with_output_types(int)
         | 'Mean' >> combine.Mean.Globally())

    self.assertEqual(float, d.element_type)
    assert_that(d, equal_to([2.0]))
    self.p.run()

  def test_mean_globally_pipeline_checking_violated(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'C' >> beam.Create(['test']).with_output_types(str)
       | 'Mean' >> combine.Mean.Globally())

    if sys.version_info[0] >= 3:
      expected_msg = \
        "Type hint violation for 'CombinePerKey': " \
        "requires Tuple[TypeVariable[K], Union[float, int]] " \
        "but got Tuple[None, str] for element"
    else:
      expected_msg = \
        "Type hint violation for 'CombinePerKey': " \
        "requires Tuple[TypeVariable[K], Union[float, int, long]] " \
        "but got Tuple[None, str] for element"

    self.assertEqual(expected_msg, e.exception.args[0])

  def test_mean_globally_runtime_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | 'C' >> beam.Create(range(5)).with_output_types(int)
         | 'Mean' >> combine.Mean.Globally())

    self.assertEqual(float, d.element_type)
    assert_that(d, equal_to([2.0]))
    self.p.run()

  def test_mean_globally_runtime_checking_violated(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'C' >> beam.Create(['t', 'e', 's', 't']).with_output_types(str)
       | 'Mean' >> combine.Mean.Globally())
      self.p.run()
      self.assertEqual("Runtime type violation detected for transform input "
                       "when executing ParDoFlatMap(Combine): Tuple[Any, "
                       "Iterable[Union[int, float]]] hint type-constraint "
                       "violated. The type of element #1 in the passed tuple "
                       "is incorrect. Iterable[Union[int, float]] hint "
                       "type-constraint violated. The type of element #0 in "
                       "the passed Iterable is incorrect: Union[int, float] "
                       "type-constraint violated. Expected an instance of one "
                       "of: ('int', 'float'), received str instead.",
                       e.exception.args[0])

  def test_mean_per_key_pipeline_checking_satisfied(self):
    d = (self.p
         | beam.Create(range(5)).with_output_types(int)
         | ('EvenGroup' >> beam.Map(lambda x: (not x % 2, x))
            .with_output_types(typing.Tuple[bool, int]))
         | 'EvenMean' >> combine.Mean.PerKey())

    self.assertCompatible(typing.Tuple[bool, float], d.element_type)
    assert_that(d, equal_to([(False, 2.0), (True, 2.0)]))
    self.p.run()

  def test_mean_per_key_pipeline_checking_violated(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create(map(str, range(5))).with_output_types(str)
       | ('UpperPair' >> beam.Map(lambda x: (x.upper(), x))
          .with_output_types(typing.Tuple[str, str]))
       | 'EvenMean' >> combine.Mean.PerKey())
      self.p.run()

    if sys.version_info[0] >= 3:
      expected_msg = \
        "Type hint violation for 'CombinePerKey(MeanCombineFn)': " \
        "requires Tuple[TypeVariable[K], Union[float, int]] " \
        "but got Tuple[str, str] for element"
    else:
      expected_msg = \
        "Type hint violation for 'CombinePerKey(MeanCombineFn)': " \
        "requires Tuple[TypeVariable[K], Union[float, int, long]] " \
        "but got Tuple[str, str] for element"

    self.assertEqual(expected_msg, e.exception.args[0])

  def test_mean_per_key_runtime_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create(range(5)).with_output_types(int)
         | ('OddGroup' >> beam.Map(lambda x: (bool(x % 2), x))
            .with_output_types(typing.Tuple[bool, int]))
         | 'OddMean' >> combine.Mean.PerKey())

    self.assertCompatible(typing.Tuple[bool, float], d.element_type)
    assert_that(d, equal_to([(False, 2.0), (True, 2.0)]))
    self.p.run()

  def test_mean_per_key_runtime_checking_violated(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create(range(5)).with_output_types(int)
       | ('OddGroup' >> beam.Map(lambda x: (x, str(bool(x % 2))))
          .with_output_types(typing.Tuple[int, str]))
       | 'OddMean' >> combine.Mean.PerKey())
      self.p.run()

    if sys.version_info[0] >= 3:
      expected_msg = \
        "Runtime type violation detected within " \
        "OddMean/CombinePerKey(MeanCombineFn): " \
        "Type-hint for argument: 'element' violated: " \
        "Union[float, int] type-constraint violated. " \
        "Expected an instance of one of: ('float', 'int'), " \
        "received str instead"
    else:
      expected_msg = \
        "Runtime type violation detected within " \
        "OddMean/CombinePerKey(MeanCombineFn): " \
        "Type-hint for argument: 'element' violated: " \
        "Union[float, int, long] type-constraint violated. " \
        "Expected an instance of one of: ('float', 'int', 'long'), " \
        "received str instead"

    self.assertStartswith(e.exception.args[0], expected_msg)

  def test_count_globally_pipeline_type_checking_satisfied(self):
    d = (self.p
         | 'P' >> beam.Create(range(5)).with_output_types(int)
         | 'CountInt' >> combine.Count.Globally())

    self.assertEqual(int, d.element_type)
    assert_that(d, equal_to([5]))
    self.p.run()

  def test_count_globally_runtime_type_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | 'P' >> beam.Create(range(5)).with_output_types(int)
         | 'CountInt' >> combine.Count.Globally())

    self.assertEqual(int, d.element_type)
    assert_that(d, equal_to([5]))
    self.p.run()

  def test_count_perkey_pipeline_type_checking_satisfied(self):
    d = (self.p
         | beam.Create(range(5)).with_output_types(int)
         | ('EvenGroup' >> beam.Map(lambda x: (not x % 2, x))
            .with_output_types(typing.Tuple[bool, int]))
         | 'CountInt' >> combine.Count.PerKey())

    self.assertCompatible(typing.Tuple[bool, int], d.element_type)
    assert_that(d, equal_to([(False, 2), (True, 3)]))
    self.p.run()

  def test_count_perkey_pipeline_type_checking_violated(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create(range(5)).with_output_types(int)
       | 'CountInt' >> combine.Count.PerKey())

    self.assertEqual(
        "Type hint violation for 'CombinePerKey(CountCombineFn)': "
        "requires Tuple[TypeVariable[K], Any] "
        "but got {} for element".format(int),
        e.exception.args[0])

  def test_count_perkey_runtime_type_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create(['t', 'e', 's', 't']).with_output_types(str)
         | 'DupKey' >> beam.Map(lambda x: (x, x))
         .with_output_types(typing.Tuple[str, str])
         | 'CountDups' >> combine.Count.PerKey())

    self.assertCompatible(typing.Tuple[str, int], d.element_type)
    assert_that(d, equal_to([('e', 1), ('s', 1), ('t', 2)]))
    self.p.run()

  def test_count_perelement_pipeline_type_checking_satisfied(self):
    d = (self.p
         | beam.Create([1, 1, 2, 3]).with_output_types(int)
         | 'CountElems' >> combine.Count.PerElement())

    self.assertCompatible(typing.Tuple[int, int], d.element_type)
    assert_that(d, equal_to([(1, 2), (2, 1), (3, 1)]))
    self.p.run()

  def test_count_perelement_pipeline_type_checking_violated(self):
    self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED'

    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | 'f' >> beam.Create([1, 1, 2, 3])
       | 'CountElems' >> combine.Count.PerElement())

    self.assertEqual('Pipeline type checking is enabled, however no output '
                     'type-hint was found for the PTransform '
                     'Create(f)',
                     e.exception.args[0])

  def test_count_perelement_runtime_type_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create([True, True, False, True, True])
         .with_output_types(bool)
         | 'CountElems' >> combine.Count.PerElement())

    self.assertCompatible(typing.Tuple[bool, int], d.element_type)
    assert_that(d, equal_to([(False, 1), (True, 4)]))
    self.p.run()

  def test_top_of_pipeline_checking_satisfied(self):
    d = (self.p
         | beam.Create(range(5, 11)).with_output_types(int)
         | 'Top 3' >> combine.Top.Of(3))

    self.assertCompatible(typing.Iterable[int],
                          d.element_type)
    assert_that(d, equal_to([[10, 9, 8]]))
    self.p.run()

  def test_top_of_runtime_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create(list('testing')).with_output_types(str)
         | 'AciiTop' >> combine.Top.Of(3))

    self.assertCompatible(typing.Iterable[str], d.element_type)
    assert_that(d, equal_to([['t', 't', 's']]))
    self.p.run()

  def test_per_key_pipeline_checking_violated(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create(range(100)).with_output_types(int)
       | 'Num + 1' >> beam.Map(lambda x: x + 1).with_output_types(int)
       | 'TopMod' >> combine.Top.PerKey(1))

    self.assertEqual(
        "Type hint violation for 'CombinePerKey(TopCombineFn)': "
        "requires Tuple[TypeVariable[K], TypeVariable[T]] "
        "but got {} for element".format(int),
        e.exception.args[0])

  def test_per_key_pipeline_checking_satisfied(self):
    d = (self.p
         | beam.Create(range(100)).with_output_types(int)
         | ('GroupMod 3' >> beam.Map(lambda x: (x % 3, x))
            .with_output_types(typing.Tuple[int, int]))
         | 'TopMod' >> combine.Top.PerKey(1))

    self.assertCompatible(typing.Tuple[int, typing.Iterable[int]],
                          d.element_type)
    assert_that(d, equal_to([(0, [99]), (1, [97]), (2, [98])]))
    self.p.run()

  def test_per_key_runtime_checking_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create(range(21))
         | ('GroupMod 3' >> beam.Map(lambda x: (x % 3, x))
            .with_output_types(typing.Tuple[int, int]))
         | 'TopMod' >> combine.Top.PerKey(1))

    self.assertCompatible(typing.Tuple[int, typing.Iterable[int]],
                          d.element_type)
    assert_that(d, equal_to([(0, [18]), (1, [19]), (2, [20])]))
    self.p.run()

  def test_sample_globally_pipeline_satisfied(self):
    d = (self.p
         | beam.Create([2, 2, 3, 3]).with_output_types(int)
         | 'Sample' >> combine.Sample.FixedSizeGlobally(3))

    self.assertCompatible(typing.Iterable[int], d.element_type)

    def matcher(expected_len):
      def match(actual):
        equal_to([expected_len])([len(actual[0])])
      return match
    assert_that(d, matcher(3))
    self.p.run()

  def test_sample_globally_runtime_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create([2, 2, 3, 3]).with_output_types(int)
         | 'Sample' >> combine.Sample.FixedSizeGlobally(2))

    self.assertCompatible(typing.Iterable[int], d.element_type)

    def matcher(expected_len):
      def match(actual):
        equal_to([expected_len])([len(actual[0])])
      return match
    assert_that(d, matcher(2))
    self.p.run()

  def test_sample_per_key_pipeline_satisfied(self):
    d = (self.p
         | (beam.Create([(1, 2), (1, 2), (2, 3), (2, 3)])
            .with_output_types(typing.Tuple[int, int]))
         | 'Sample' >> combine.Sample.FixedSizePerKey(2))

    self.assertCompatible(typing.Tuple[int, typing.Iterable[int]],
                          d.element_type)

    def matcher(expected_len):
      def match(actual):
        for _, sample in actual:
          equal_to([expected_len])([len(sample)])
      return match
    assert_that(d, matcher(2))
    self.p.run()

  def test_sample_per_key_runtime_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | (beam.Create([(1, 2), (1, 2), (2, 3), (2, 3)])
            .with_output_types(typing.Tuple[int, int]))
         | 'Sample' >> combine.Sample.FixedSizePerKey(1))

    self.assertCompatible(typing.Tuple[int, typing.Iterable[int]],
                          d.element_type)

    def matcher(expected_len):
      def match(actual):
        for _, sample in actual:
          equal_to([expected_len])([len(sample)])
      return match
    assert_that(d, matcher(1))
    self.p.run()

  def test_to_list_pipeline_check_satisfied(self):
    d = (self.p
         | beam.Create((1, 2, 3, 4)).with_output_types(int)
         | combine.ToList())

    self.assertCompatible(typing.List[int], d.element_type)

    def matcher(expected):
      def match(actual):
        equal_to(expected)(actual[0])
      return match
    assert_that(d, matcher([1, 2, 3, 4]))
    self.p.run()

  def test_to_list_runtime_check_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | beam.Create(list('test')).with_output_types(str)
         | combine.ToList())

    self.assertCompatible(typing.List[str], d.element_type)

    def matcher(expected):
      def match(actual):
        equal_to(expected)(actual[0])
      return match
    assert_that(d, matcher(['e', 's', 't', 't']))
    self.p.run()

  def test_to_dict_pipeline_check_violated(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      (self.p
       | beam.Create([1, 2, 3, 4]).with_output_types(int)
       | combine.ToDict())

    self.assertEqual(
        "Type hint violation for 'CombinePerKey': "
        "requires "
        "Tuple[TypeVariable[K], Tuple[TypeVariable[K], TypeVariable[V]]] "
        "but got Tuple[None, int] for element",
        e.exception.args[0])

  def test_to_dict_pipeline_check_satisfied(self):
    d = (self.p
         | beam.Create(
             [(1, 2), (3, 4)]).with_output_types(typing.Tuple[int, int])
         | combine.ToDict())

    self.assertCompatible(typing.Dict[int, int], d.element_type)
    assert_that(d, equal_to([{1: 2, 3: 4}]))
    self.p.run()

  def test_to_dict_runtime_check_satisfied(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    d = (self.p
         | (beam.Create([('1', 2), ('3', 4)])
            .with_output_types(typing.Tuple[str, int]))
         | combine.ToDict())

    self.assertCompatible(typing.Dict[str, int], d.element_type)
    assert_that(d, equal_to([{'1': 2, '3': 4}]))
    self.p.run()

  def test_runtime_type_check_python_type_error(self):
    self.p._options.view_as(TypeOptions).runtime_type_check = True

    with self.assertRaises(TypeError) as e:
      (self.p
       | beam.Create([1, 2, 3]).with_output_types(int)
       | 'Len' >> beam.Map(lambda x: len(x)).with_output_types(int))
      self.p.run()

    # Our special type-checking related TypeError shouldn't have been raised.
    # Instead the above pipeline should have triggered a regular Python runtime
    # TypeError.
    self.assertEqual("object of type 'int' has no len() [while running 'Len']",
                     e.exception.args[0])
    self.assertFalse(isinstance(e, typehints.TypeCheckError))

  def test_pardo_type_inference(self):
    self.assertEqual(int,
                     beam.Filter(lambda x: False).infer_output_type(int))
    self.assertEqual(typehints.Tuple[str, int],
                     beam.Map(lambda x: (x, 1)).infer_output_type(str))

  def test_gbk_type_inference(self):
    self.assertEqual(
        typehints.Tuple[str, typehints.Iterable[int]],
        _GroupByKeyOnly().infer_output_type(typehints.KV[str, int]))

  def test_pipeline_inference(self):
    created = self.p | beam.Create(['a', 'b', 'c'])
    mapped = created | 'pair with 1' >> beam.Map(lambda x: (x, 1))
    grouped = mapped | beam.GroupByKey()
    self.assertEqual(str, created.element_type)
    self.assertEqual(typehints.KV[str, int], mapped.element_type)
    self.assertEqual(typehints.KV[str, typehints.Iterable[int]],
                     grouped.element_type)

  def test_inferred_bad_kv_type(self):
    with self.assertRaises(typehints.TypeCheckError) as e:
      _ = (self.p
           | beam.Create(['a', 'b', 'c'])
           | 'Ungroupable' >> beam.Map(lambda x: (x, 0, 1.0))
           | beam.GroupByKey())

    self.assertEqual('Input type hint violation at GroupByKey: '
                     'expected Tuple[TypeVariable[K], TypeVariable[V]], '
                     'got Tuple[str, int, float]',
                     e.exception.args[0])

  def test_type_inference_command_line_flag_toggle(self):
    self.p._options.view_as(TypeOptions).pipeline_type_check = False
    x = self.p | 'C1' >> beam.Create([1, 2, 3, 4])
    self.assertIsNone(x.element_type)

    self.p._options.view_as(TypeOptions).pipeline_type_check = True
    x = self.p | 'C2' >> beam.Create([1, 2, 3, 4])
    self.assertEqual(int, x.element_type)

  def test_eager_execution(self):
    doubled = [1, 2, 3, 4] | beam.Map(lambda x: 2 * x)
    self.assertEqual([2, 4, 6, 8], doubled)

  def test_eager_execution_tagged_outputs(self):
    result = [1, 2, 3, 4] | beam.Map(
        lambda x: pvalue.TaggedOutput('bar', 2 * x)).with_outputs('bar')
    self.assertEqual([2, 4, 6, 8], result.bar)
    with self.assertRaises(KeyError,
                           msg='Tag \'foo\' is not a defined output tag'):
      result.foo


class TestPTransformFn(TypeHintTestCase):

  def test_type_checking_fail(self):
    @beam.ptransform_fn
    def MyTransform(pcoll):
      return pcoll | beam.ParDo(lambda x: [x]).with_output_types(str)

    p = TestPipeline()
    with self.assertRaisesRegex(beam.typehints.TypeCheckError,
                                r'expected.*int.*got.*str'):
      _ = (p
           | beam.Create([1, 2])
           | MyTransform().with_output_types(int))

  def test_type_checking_success(self):
    @beam.ptransform_fn
    def MyTransform(pcoll):
      return pcoll | beam.ParDo(lambda x: [x]).with_output_types(int)

    p = TestPipeline()
    _ = (p
         | beam.Create([1, 2])
         | MyTransform().with_output_types(int))
    p.run()

  def test_type_hints_arg(self):
    # Tests passing type hints via the magic 'type_hints' argument name.
    @beam.ptransform_fn
    def MyTransform(pcoll, type_hints, test_arg):
      self.assertEqual(test_arg, 'test')
      return (pcoll
              | beam.ParDo(lambda x: [x]).with_output_types(
                  type_hints.output_types[0][0]))

    p = TestPipeline()
    _ = (p
         | beam.Create([1, 2])
         | MyTransform('test').with_output_types(int))
    p.run()


if __name__ == '__main__':
  unittest.main()
