#
# 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 type-hint objects and decorators."""

from __future__ import absolute_import

import sys
import typing
import unittest

import apache_beam as beam
from apache_beam import pvalue
from apache_beam import typehints
from apache_beam.options.pipeline_options import OptionsContext
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 WithTypeHints
from apache_beam.typehints import decorators
from apache_beam.typehints.decorators import get_signature

# These test often construct a pipeline as value | PTransform to test side
# effects (e.g. errors).
# pylint: disable=expression-not-assigned

decorators._enable_from_callable = True


class MainInputTest(unittest.TestCase):

  def test_bad_main_input(self):
    @typehints.with_input_types(str, int)
    def repeat(s, times):
      return s * times
    with self.assertRaises(typehints.TypeCheckError):
      [1, 2, 3] | beam.Map(repeat, 3)

  def test_non_function(self):
    result = ['a', 'bb', 'c'] | beam.Map(str.upper)
    self.assertEqual(['A', 'BB', 'C'], sorted(result))

    result = ['xa', 'bbx', 'xcx'] | beam.Map(str.strip, 'x')
    self.assertEqual(['a', 'bb', 'c'], sorted(result))

    result = ['1', '10', '100'] | beam.Map(int)
    self.assertEqual([1, 10, 100], sorted(result))

    result = ['1', '10', '100'] | beam.Map(int, 16)
    self.assertEqual([1, 16, 256], sorted(result))

  @unittest.skipIf(
      sys.version_info.major >= 3 and sys.version_info < (3, 7, 0),
      'Function signatures for builtins are not available in Python 3 before '
      'version 3.7.')
  def test_non_function_fails(self):
    with self.assertRaises(typehints.TypeCheckError):
      [1, 2, 3] | beam.Map(str.upper)

  def test_loose_bounds(self):
    @typehints.with_input_types(typing.Union[int, float])
    @typehints.with_output_types(str)
    def format_number(x):
      return '%g' % x
    result = [1, 2, 3] | beam.Map(format_number)
    self.assertEqual(['1', '2', '3'], sorted(result))

  def test_typed_dofn_class(self):
    @typehints.with_input_types(int)
    @typehints.with_output_types(str)
    class MyDoFn(beam.DoFn):
      def process(self, element):
        return [str(element)]

    result = [1, 2, 3] | beam.ParDo(MyDoFn())
    self.assertEqual(['1', '2', '3'], sorted(result))

    with self.assertRaisesRegexp(typehints.TypeCheckError,
                                 r'requires.*int.*got.*str'):
      ['a', 'b', 'c'] | beam.ParDo(MyDoFn())

    with self.assertRaisesRegexp(typehints.TypeCheckError,
                                 r'requires.*int.*got.*str'):
      [1, 2, 3] | (beam.ParDo(MyDoFn()) | 'again' >> beam.ParDo(MyDoFn()))

  @unittest.skip('BEAM-7981: Iterable in output type should not be removed.')
  def test_typed_callable_iterable_output(self):
    # TODO(BEAM-7981): 2.7 and 3.x both erroneously strip the Iterable, but the
    #   test only fails in 3.x.
    @typehints.with_input_types(int)
    @typehints.with_output_types(typehints.Iterable[str])
    def do_fn(element):
      return [[str(element)] * 2]

    result = [1, 2] | beam.ParDo(do_fn)
    self.assertEqual([['1', '1'], ['2', '2']], sorted(result))

  def test_typed_dofn_instance(self):
    class MyDoFn(beam.DoFn):
      def process(self, element):
        return [str(element)]
    my_do_fn = MyDoFn().with_input_types(int).with_output_types(str)

    result = [1, 2, 3] | beam.ParDo(my_do_fn)
    self.assertEqual(['1', '2', '3'], sorted(result))

    with self.assertRaises(typehints.TypeCheckError):
      ['a', 'b', 'c'] | beam.ParDo(my_do_fn)

    with self.assertRaises(typehints.TypeCheckError):
      [1, 2, 3] | (beam.ParDo(my_do_fn) | 'again' >> beam.ParDo(my_do_fn))

  def test_filter_type_hint(self):
    @typehints.with_input_types(int)
    def filter_fn(data):
      return data % 2

    self.assertEqual([1, 3], [1, 2, 3] | beam.Filter(filter_fn))


class NativeTypesTest(unittest.TestCase):

  def test_good_main_input(self):
    @typehints.with_input_types(typing.Tuple[str, int])
    def munge(s_i):
      (s, i) = s_i
      return (s + 's', i * 2)
    result = [('apple', 5), ('pear', 3)] | beam.Map(munge)
    self.assertEqual([('apples', 10), ('pears', 6)], sorted(result))

  def test_bad_main_input(self):
    @typehints.with_input_types(typing.Tuple[str, str])
    def munge(s_i):
      (s, i) = s_i
      return (s + 's', i * 2)
    with self.assertRaises(typehints.TypeCheckError):
      [('apple', 5), ('pear', 3)] | beam.Map(munge)

  def test_bad_main_output(self):
    @typehints.with_input_types(typing.Tuple[int, int])
    @typehints.with_output_types(typing.Tuple[str, str])
    def munge(a_b):
      (a, b) = a_b
      return (str(a), str(b))
    with self.assertRaises(typehints.TypeCheckError):
      [(5, 4), (3, 2)] | beam.Map(munge) | 'Again' >> beam.Map(munge)


class SideInputTest(unittest.TestCase):

  def _run_repeat_test(self, repeat):
    self._run_repeat_test_good(repeat)
    self._run_repeat_test_bad(repeat)

  @OptionsContext(pipeline_type_check=True)
  def _run_repeat_test_good(self, repeat):
    # As a positional argument.
    result = ['a', 'bb', 'c'] | beam.Map(repeat, 3)
    self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result))

    # As a keyword argument.
    result = ['a', 'bb', 'c'] | beam.Map(repeat, times=3)
    self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result))

  def _run_repeat_test_bad(self, repeat):
    # Various mismatches.
    with self.assertRaises(typehints.TypeCheckError):
      ['a', 'bb', 'c'] | beam.Map(repeat, 'z')
    with self.assertRaises(typehints.TypeCheckError):
      ['a', 'bb', 'c'] | beam.Map(repeat, times='z')
    with self.assertRaises(typehints.TypeCheckError):
      ['a', 'bb', 'c'] | beam.Map(repeat, 3, 4)
    if all(param.default == param.empty
           for param in get_signature(repeat).parameters.values()):
      with self.assertRaisesRegexp(typehints.TypeCheckError,
                                   r'(takes exactly|missing a required)'):
        ['a', 'bb', 'c'] | beam.Map(repeat)

  def test_basic_side_input_hint(self):
    @typehints.with_input_types(str, int)
    def repeat(s, times):
      return s * times
    self._run_repeat_test(repeat)

  def test_keyword_side_input_hint(self):
    @typehints.with_input_types(str, times=int)
    def repeat(s, times):
      return s * times
    self._run_repeat_test(repeat)

  def test_default_typed_hint(self):
    @typehints.with_input_types(str, int)
    def repeat(s, times=3):
      return s * times
    self._run_repeat_test(repeat)

  def test_default_untyped_hint(self):
    @typehints.with_input_types(str)
    def repeat(s, times=3):
      return s * times
    # No type checking on default arg.
    self._run_repeat_test_good(repeat)

  @OptionsContext(pipeline_type_check=True)
  def test_varargs_side_input_hint(self):
    @typehints.with_input_types(str, int)
    def repeat(s, *times):
      return s * times[0]

    result = ['a', 'bb', 'c'] | beam.Map(repeat, 3)
    self.assertEqual(['aaa', 'bbbbbb', 'ccc'], sorted(result))

    if sys.version_info >= (3,):
      with self.assertRaisesRegexp(
          typehints.TypeCheckError,
          r'requires Tuple\[int, ...\] but got Tuple\[str, ...\]'):
        ['a', 'bb', 'c'] | beam.Map(repeat, 'z')

  def test_var_positional_only_side_input_hint(self):
    # Test that a lambda that accepts only a VAR_POSITIONAL can accept
    # side-inputs.
    # TODO(BEAM-8247): There's a bug with trivial_inference inferring the output
    #   type when side-inputs are used (their type hints are not passed). Remove
    #   with_output_types(...) when this bug is fixed.
    result = (['a', 'b', 'c']
              | beam.Map(lambda *args: args, 5).with_input_types(int, str)
              .with_output_types(typehints.Tuple[str, int]))
    self.assertEqual([('a', 5), ('b', 5), ('c', 5)], sorted(result))

    # Type hint order doesn't matter for VAR_POSITIONAL.
    result = (['a', 'b', 'c']
              | beam.Map(lambda *args: args, 5).with_input_types(int, str)
              .with_output_types(typehints.Tuple[str, int]))
    self.assertEqual([('a', 5), ('b', 5), ('c', 5)], sorted(result))

    if sys.version_info >= (3,):
      with self.assertRaisesRegexp(
          typehints.TypeCheckError,
          r'requires Tuple\[Union\[int, str\], ...\] but got '
          r'Tuple\[Union\[float, int\], ...\]'):
        _ = [1.2] | beam.Map(lambda *_: 'a', 5).with_input_types(int, str)

  def test_var_keyword_side_input_hint(self):
    # Test that a lambda that accepts a VAR_KEYWORD can accept
    # side-inputs.
    result = (['a', 'b', 'c']
              | beam.Map(lambda e, **kwargs: (e, kwargs), kw=5)
              .with_input_types(str, ignored=int))
    self.assertEqual([('a', {'kw': 5}), ('b', {'kw': 5}), ('c', {'kw': 5})],
                     sorted(result))

    if sys.version_info >= (3,):
      with self.assertRaisesRegexp(
          typehints.TypeCheckError,
          r'requires Dict\[str, str\] but got Dict\[str, int\]'):
        _ = (['a', 'b', 'c']
             | beam.Map(lambda e, **_: 'a', kw=5)
             .with_input_types(str, ignored=str))

  def test_deferred_side_inputs(self):
    @typehints.with_input_types(str, int)
    def repeat(s, times):
      return s * times
    with TestPipeline() as p:
      main_input = p | beam.Create(['a', 'bb', 'c'])
      side_input = p | 'side' >> beam.Create([3])
      result = main_input | beam.Map(repeat, pvalue.AsSingleton(side_input))
      assert_that(result, equal_to(['aaa', 'bbbbbb', 'ccc']))

    bad_side_input = p | 'bad_side' >> beam.Create(['z'])
    with self.assertRaises(typehints.TypeCheckError):
      main_input | 'bis' >> beam.Map(repeat, pvalue.AsSingleton(bad_side_input))

  def test_deferred_side_input_iterable(self):
    @typehints.with_input_types(str, typing.Iterable[str])
    def concat(glue, items):
      return glue.join(sorted(items))
    with TestPipeline() as p:
      main_input = p | beam.Create(['a', 'bb', 'c'])
      side_input = p | 'side' >> beam.Create(['x', 'y', 'z'])
      result = main_input | beam.Map(concat, pvalue.AsIter(side_input))
      assert_that(result, equal_to(['xayaz', 'xbbybbz', 'xcycz']))

    bad_side_input = p | 'bad_side' >> beam.Create([1, 2, 3])
    with self.assertRaises(typehints.TypeCheckError):
      main_input | 'fail' >> beam.Map(concat, pvalue.AsIter(bad_side_input))


class CustomTransformTest(unittest.TestCase):

  class CustomTransform(beam.PTransform):

    def _extract_input_pvalues(self, pvalueish):
      return pvalueish, (pvalueish['in0'], pvalueish['in1'])

    def expand(self, pvalueish):
      return {'out0': pvalueish['in0'], 'out1': pvalueish['in1']}

    # TODO(robertwb): (typecheck) Make these the default?
    def with_input_types(self, *args, **kwargs):
      return WithTypeHints.with_input_types(self, *args, **kwargs)

    def with_output_types(self, *args, **kwargs):
      return WithTypeHints.with_output_types(self, *args, **kwargs)

  test_input = {'in0': ['a', 'b', 'c'], 'in1': [1, 2, 3]}

  def check_output(self, result):
    self.assertEqual(['a', 'b', 'c'], sorted(result['out0']))
    self.assertEqual([1, 2, 3], sorted(result['out1']))

  def test_custom_transform(self):
    self.check_output(self.test_input | self.CustomTransform())

  def test_keyword_type_hints(self):
    self.check_output(
        self.test_input | self.CustomTransform().with_input_types(
            in0=str, in1=int))
    self.check_output(
        self.test_input | self.CustomTransform().with_input_types(in0=str))
    self.check_output(
        self.test_input | self.CustomTransform().with_output_types(
            out0=str, out1=int))
    with self.assertRaises(typehints.TypeCheckError):
      self.test_input | self.CustomTransform().with_input_types(in0=int)
    with self.assertRaises(typehints.TypeCheckError):
      self.test_input | self.CustomTransform().with_output_types(out0=int)

  def test_flat_type_hint(self):
    # Type hint is applied to both.
    ({'in0': ['a', 'b', 'c'], 'in1': ['x', 'y', 'z']}
     | self.CustomTransform().with_input_types(str))
    with self.assertRaises(typehints.TypeCheckError):
      self.test_input | self.CustomTransform().with_input_types(str)
    with self.assertRaises(typehints.TypeCheckError):
      self.test_input | self.CustomTransform().with_input_types(int)
    with self.assertRaises(typehints.TypeCheckError):
      self.test_input | self.CustomTransform().with_output_types(int)


class AnnotationsTest(unittest.TestCase):

  def test_pardo_wrapper_builtin_method(self):
    th = beam.ParDo(str.strip).get_type_hints()
    if sys.version_info < (3, 7):
      self.assertEqual(th.input_types, ((str,), {}))
    else:
      # Python 3.7+ has annotations for CPython builtins
      # (_MethodDescriptorType).
      self.assertEqual(th.input_types, ((str, typehints.Any), {}))
    self.assertEqual(th.output_types, ((typehints.Any,), {}))

  def test_pardo_wrapper_builtin_type(self):
    th = beam.ParDo(list).get_type_hints()
    if sys.version_info < (3, 7):
      self.assertEqual(th.input_types, (
          (typehints.Any, typehints.decorators._ANY_VAR_POSITIONAL),
          {'__unknown__keywords': typehints.decorators._ANY_VAR_KEYWORD}))
    else:
      # Python 3.7+ supports signatures for builtins like 'list'.
      self.assertEqual(th.input_types, ((typehints.Any, ), {}))

    self.assertEqual(th.output_types, ((typehints.Any,), {}))

  def test_pardo_wrapper_builtin_func(self):
    th = beam.ParDo(len).get_type_hints()
    self.assertIsNone(th.input_types)
    self.assertIsNone(th.output_types)


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