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

"""Tests for apache_beam.typehints.trivial_inference."""

from __future__ import absolute_import

import sys
import unittest

from apache_beam.typehints import trivial_inference
from apache_beam.typehints import typehints

global_int = 1


class TrivialInferenceTest(unittest.TestCase):

  def assertReturnType(self, expected, f, inputs=(), depth=5):
    self.assertEqual(
        expected,
        trivial_inference.infer_return_type(f, inputs, debug=True, depth=depth))

  def testIdentity(self):
    self.assertReturnType(int, lambda x: x, [int])

  def testIndexing(self):
    self.assertReturnType(int, lambda x: x[0], [typehints.Tuple[int, str]])
    self.assertReturnType(str, lambda x: x[1], [typehints.Tuple[int, str]])
    self.assertReturnType(str, lambda x: x[1], [typehints.List[str]])

  def testTuples(self):
    self.assertReturnType(
        typehints.Tuple[typehints.Tuple[()], int], lambda x: ((), x), [int])
    self.assertReturnType(
        typehints.Tuple[str, int, float], lambda x: (x, 0, 1.0), [str])

  def testGetItem(self):
    def reverse(ab):
      return ab[-1], ab[0]
    self.assertReturnType(
        typehints.Tuple[typehints.Any, typehints.Any], reverse, [typehints.Any])
    self.assertReturnType(
        typehints.Tuple[int, float], reverse, [typehints.Tuple[float, int]])
    self.assertReturnType(
        typehints.Tuple[int, str], reverse, [typehints.Tuple[str, float, int]])
    self.assertReturnType(
        typehints.Tuple[int, int], reverse, [typehints.List[int]])

  def testGetItemSlice(self):
    self.assertReturnType(
        typehints.List[int], lambda v: v[::-1], [typehints.List[int]])
    self.assertReturnType(
        typehints.Tuple[int], lambda v: v[::-1], [typehints.Tuple[int]])
    self.assertReturnType(str, lambda v: v[::-1], [str])
    self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any])
    self.assertReturnType(typehints.Any, lambda v: v[::-1], [object])

  def testUnpack(self):
    def reverse(a_b):
      (a, b) = a_b
      return b, a
    any_tuple = typehints.Tuple[typehints.Any, typehints.Any]
    self.assertReturnType(
        typehints.Tuple[int, float], reverse, [typehints.Tuple[float, int]])
    self.assertReturnType(
        typehints.Tuple[int, int], reverse, [typehints.Tuple[int, ...]])
    self.assertReturnType(
        typehints.Tuple[int, int], reverse, [typehints.List[int]])
    self.assertReturnType(
        typehints.Tuple[typehints.Union[int, float, str],
                        typehints.Union[int, float, str]],
        reverse, [typehints.Tuple[int, float, str]])
    self.assertReturnType(any_tuple, reverse, [typehints.Any])

    self.assertReturnType(typehints.Tuple[int, float],
                          reverse, [trivial_inference.Const((1.0, 1))])
    self.assertReturnType(any_tuple,
                          reverse, [trivial_inference.Const((1, 2, 3))])

  def testNoneReturn(self):
    def func(a):
      if a == 5:
        return a
      return None
    self.assertReturnType(typehints.Union[int, type(None)], func, [int])

  def testSimpleList(self):
    self.assertReturnType(
        typehints.List[int],
        lambda xs: [1, 2],
        [typehints.Tuple[int, ...]])

    self.assertReturnType(
        typehints.List[typehints.Any],
        lambda xs: list(xs), # List is a disallowed builtin
        [typehints.Tuple[int, ...]])

  def testListComprehension(self):
    self.assertReturnType(
        typehints.List[int],
        lambda xs: [x for x in xs],
        [typehints.Tuple[int, ...]])

  def testTupleListComprehension(self):
    self.assertReturnType(
        typehints.List[int],
        lambda xs: [x for x in xs],
        [typehints.Tuple[int, int, int]])
    self.assertReturnType(
        typehints.List[typehints.Union[int, float]],
        lambda xs: [x for x in xs],
        [typehints.Tuple[int, float]])
    if sys.version_info[:2] == (3, 5):
      # A better result requires implementing the MAKE_CLOSURE opcode.
      expected = typehints.Any
    else:
      expected = typehints.List[typehints.Tuple[str, int]]
    self.assertReturnType(
        expected,
        lambda kvs: [(kvs[0], v) for v in kvs[1]],
        [typehints.Tuple[str, typehints.Iterable[int]]])
    self.assertReturnType(
        typehints.List[typehints.Tuple[str, typehints.Union[str, int], int]],
        lambda L: [(a, a or b, b) for a, b in L],
        [typehints.Iterable[typehints.Tuple[str, int]]])

  def testGenerator(self):

    def foo(x, y):
      yield x
      yield y

    self.assertReturnType(typehints.Iterable[int], foo, [int, int])
    self.assertReturnType(
        typehints.Iterable[typehints.Union[int, float]], foo, [int, float])

  def testGeneratorComprehension(self):
    self.assertReturnType(
        typehints.Iterable[int],
        lambda xs: (x for x in xs),
        [typehints.Tuple[int, ...]])

  def testBinOp(self):
    self.assertReturnType(int, lambda a, b: a + b, [int, int])
    self.assertReturnType(
        typehints.Any, lambda a, b: a + b, [int, typehints.Any])
    self.assertReturnType(
        typehints.List[typehints.Union[int, str]], lambda a, b: a + b,
        [typehints.List[int], typehints.List[str]])

  def testCall(self):
    f = lambda x, *args: x
    self.assertReturnType(
        typehints.Tuple[int, float], lambda: (f(1), f(2.0, 3)))
    # We could do better here, but this is at least correct.
    self.assertReturnType(
        typehints.Tuple[int, typehints.Any], lambda: (1, f(x=1.0)))

  def testClosure(self):
    x = 1
    y = 1.0
    self.assertReturnType(typehints.Tuple[int, float], lambda: (x, y))

  def testGlobals(self):
    self.assertReturnType(int, lambda: global_int)

  def testBuiltins(self):
    self.assertReturnType(int, lambda x: len(x), [typehints.Any])

  def testGetAttr(self):
    self.assertReturnType(
        typehints.Tuple[str, typehints.Any],
        lambda: (typehints.__doc__, typehints.fake))

  def testMethod(self):

    class A(object):
      def m(self, x):
        return x

    self.assertReturnType(int, lambda: A().m(3))
    self.assertReturnType(float, lambda: A.m(A(), 3.0))

  def testAlwaysReturnsEarly(self):

    def some_fn(v):
      if v:
        return 1
      return 2

    self.assertReturnType(int, some_fn)

  def testDict(self):
    self.assertReturnType(
        typehints.Dict[typehints.Any, typehints.Any], lambda: {})

  def testDictComprehension(self):
    fields = []
    if sys.version_info >= (3, 6):
      expected_type = typehints.Dict[typehints.Any, typehints.Any]
    else:
      # For Python 2, just ensure it doesn't crash.
      expected_type = typehints.Any
    self.assertReturnType(
        expected_type,
        lambda row: {f: row[f] for f in fields}, [typehints.Any])

  def testDictComprehensionSimple(self):
    self.assertReturnType(
        typehints.Dict[str, int],
        lambda _list: {'a': 1 for _ in _list}, [])

  def testDepthFunction(self):
    def f(i):
      return i
    self.assertReturnType(typehints.Any, lambda i: f(i), [int], depth=0)
    self.assertReturnType(int, lambda i: f(i), [int], depth=1)

  def testDepthMethod(self):
    class A(object):
      def m(self, x):
        return x

    self.assertReturnType(typehints.Any, lambda: A().m(3), depth=0)
    self.assertReturnType(int, lambda: A().m(3), depth=1)

    self.assertReturnType(typehints.Any, lambda: A.m(A(), 3.0), depth=0)
    self.assertReturnType(float, lambda: A.m(A(), 3.0), depth=1)

  def testBuildTupleUnpackWithCall(self):
    # Lambda uses BUILD_TUPLE_UNPACK_WITH_CALL opcode in Python 3.6, 3.7.
    def fn(x1, x2, *unused_args):
      return x1, x2

    self.assertReturnType(typehints.Tuple[str, float],
                          lambda x1, x2, _list: fn(x1, x2, *_list),
                          [str, float, typehints.List[int]])
    # No *args
    self.assertReturnType(typehints.Tuple[str, typehints.List[int]],
                          lambda x1, x2, _list: fn(x1, x2, *_list),
                          [str, typehints.List[int]])

  @unittest.skipIf(sys.version_info < (3, 6), 'CALL_FUNCTION_EX is new in 3.6')
  def testCallFunctionEx(self):
    # Test when fn arguments are built using BUiLD_LIST.
    def fn(*args):
      return args

    self.assertReturnType(typehints.List[typehints.Union[str, float]],
                          lambda x1, x2: fn(*[x1, x2]),
                          [str, float])

  @unittest.skipIf(sys.version_info < (3, 6), 'CALL_FUNCTION_EX is new in 3.6')
  def testCallFunctionExKwargs(self):
    def fn(x1, x2, **unused_kwargs):
      return x1, x2

    # Keyword args are currently unsupported for CALL_FUNCTION_EX.
    self.assertReturnType(typehints.Any,
                          lambda x1, x2, _dict: fn(x1, x2, **_dict),
                          [str, float, typehints.List[int]])


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