#

# 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 decorators module."""

# pytype: skip-file

import functools

import typing

import unittest

from apache_beam import Map

from apache_beam.pvalue import TaggedOutput

from apache_beam.typehints import Any

from apache_beam.typehints import Dict

from apache_beam.typehints import List

from apache_beam.typehints import Tuple

from apache_beam.typehints import TypeCheckError

from apache_beam.typehints import TypeVariable

from apache_beam.typehints import WithTypeHints

from apache_beam.typehints import decorators

from apache_beam.typehints import typehints

from apache_beam.typehints.native_type_compatibility import convert_to_beam_type

T = TypeVariable('T')

# Name is 'T' so it converts to a beam type with the same name.

# mypy requires that the name of the variable match, so we must ignore this.

# pylint: disable=typevar-name-mismatch

T_typing = typing.TypeVar('T') # type: ignore

class IOTypeHintsTest(unittest.TestCase):

def test_get_signature(self):

# Basic coverage only to make sure function works.

def fn(a, b=1, *c, **d):

return a, b, c, d

s = decorators.get_signature(fn)

self.assertListEqual(list(s.parameters), ['a', 'b', 'c', 'd'])

def test_get_signature_builtin(self):

s = decorators.get_signature(list)

self.assertListEqual(list(s.parameters), ['iterable'])

self.assertEqual(s.return_annotation, List[Any])

def test_from_callable_without_annotations(self):

def fn(a, b=None, *args, **kwargs):

return a, b, args, kwargs

th = decorators.IOTypeHints.from_callable(fn)

self.assertIsNone(th)

def test_from_callable_builtin(self):

th = decorators.IOTypeHints.from_callable(len)

self.assertIsNone(th)

def test_from_callable_method_descriptor(self):

# from_callable() injects an annotation in this special type of builtin.

th = decorators.IOTypeHints.from_callable(str.strip)

self.assertEqual(th.input_types, ((str, Any), {}))

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

def test_strip_iterable_not_simple_output_noop(self):

th = decorators.IOTypeHints(

input_types=None, output_types=((int, str), {}), origin=[])

th = th.strip_iterable()

self.assertEqual(((int, str), {}), th.output_types)

def _test_strip_iterable(self, before, expected_after):

th = decorators.IOTypeHints(

input_types=None, output_types=((before, ), {}), origin=[])

after = th.strip_iterable()

self.assertEqual(((expected_after, ), {}), after.output_types)

def test_with_output_types_from(self):

th = decorators.IOTypeHints(

input_types=((int), {

'foo': str

}),

output_types=((int, str), {}),

origin=[])

self.assertEqual(

th.with_output_types_from(decorators.IOTypeHints.empty()),

decorators.IOTypeHints(

input_types=((int), {

'foo': str

}), output_types=None, origin=[]))

self.assertEqual(

decorators.IOTypeHints.empty().with_output_types_from(th),

decorators.IOTypeHints(

input_types=None, output_types=((int, str), {}), origin=[]))

def test_with_input_types_from(self):

th = decorators.IOTypeHints(

input_types=((int), {

'foo': str

}),

output_types=((int, str), {}),

origin=[])

self.assertEqual(

th.with_input_types_from(decorators.IOTypeHints.empty()),

decorators.IOTypeHints(

input_types=None, output_types=((int, str), {}), origin=[]))

self.assertEqual(

decorators.IOTypeHints.empty().with_input_types_from(th),

decorators.IOTypeHints(

input_types=((int), {

'foo': str

}), output_types=None, origin=[]))

def _test_strip_iterable_fail(self, before):

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

self._test_strip_iterable(before, None)

def test_strip_iterable(self):

self._test_strip_iterable(None, None)

self._test_strip_iterable(typehints.Any, typehints.Any)

self._test_strip_iterable(T, typehints.Any)

self._test_strip_iterable(typehints.Iterable[str], str)

self._test_strip_iterable(typehints.List[str], str)

self._test_strip_iterable(typehints.Iterator[str], str)

self._test_strip_iterable(typehints.Generator[str], str)

self._test_strip_iterable(typehints.Tuple[str], str)

self._test_strip_iterable(

typehints.Tuple[str, int], typehints.Union[str, int])

self._test_strip_iterable(typehints.Tuple[str, ...], str)

self._test_strip_iterable(typehints.KV[str, int], typehints.Union[str, int])

self._test_strip_iterable(typehints.Set[str], str)

self._test_strip_iterable(typehints.FrozenSet[str], str)

self._test_strip_iterable_fail(typehints.Union[str, int])

self._test_strip_iterable_fail(typehints.Optional[str])

self._test_strip_iterable_fail(

typehints.WindowedValue[str]) # type: ignore[misc]

self._test_strip_iterable_fail(typehints.Dict[str, int])

def test_make_traceback(self):

origin = ''.join(

decorators.IOTypeHints.empty().with_input_types(str).origin)

self.assertRegex(origin, __name__)

self.assertNotRegex(origin, r'\b_make_traceback')

def test_origin(self):

th = decorators.IOTypeHints.empty()

self.assertListEqual([], th.origin)

th = th.with_input_types(str)

self.assertRegex(th.debug_str(), r'with_input_types')

th = th.with_output_types(str)

self.assertRegex(th.debug_str(), r'(?s)with_output_types.*with_input_types')

th = decorators.IOTypeHints.empty().with_output_types(str)

th2 = decorators.IOTypeHints.empty().with_input_types(int)

th = th.with_defaults(th2)

self.assertRegex(th.debug_str(), r'(?s)based on:.*\'str\'.*and:.*\'int\'')

def test_with_defaults_noop_does_not_grow_origin(self):

th = decorators.IOTypeHints.empty()

expected_id = id(th)

th = th.with_defaults(None)

self.assertEqual(expected_id, id(th))

th = th.with_defaults(decorators.IOTypeHints.empty())

self.assertEqual(expected_id, id(th))

th = th.with_input_types(str)

expected_id = id(th)

th = th.with_defaults(th)

self.assertEqual(expected_id, id(th))

th2 = th.with_output_types(int)

th = th.with_defaults(th2)

self.assertNotEqual(expected_id, id(th))

def test_from_callable(self):

def fn(

a: int,

b: str = '',

*args: Tuple[T],

foo: List[int],

**kwargs: Dict[str, str]) -> Tuple[Any, ...]:

return a, b, args, foo, kwargs

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(

th.input_types, ((int, str, Tuple[T]), {

'foo': List[int], 'kwargs': Dict[str, str]

}))

self.assertEqual(th.output_types, ((Tuple[Any, ...], ), {}))

def test_from_callable_partial_annotations(self):

def fn(a: int, b=None, *args, foo: List[int], **kwargs):

return a, b, args, foo, kwargs

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(

th.input_types,

((int, Any, Tuple[Any, ...]), {

'foo': List[int], 'kwargs': Dict[Any, Any]

}))

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

def test_from_callable_class(self):

class Class(object):

def __init__(self, unused_arg: int):

pass

th = decorators.IOTypeHints.from_callable(Class)

self.assertEqual(th.input_types, ((int, ), {}))

self.assertEqual(th.output_types, ((Class, ), {}))

def test_from_callable_method(self):

class Class(object):

def method(self, arg: T = None) -> None:

pass

th = decorators.IOTypeHints.from_callable(Class.method)

self.assertEqual(th.input_types, ((Any, T), {}))

self.assertEqual(th.output_types, ((None, ), {}))

th = decorators.IOTypeHints.from_callable(Class().method)

self.assertEqual(th.input_types, ((T, ), {}))

self.assertEqual(th.output_types, ((None, ), {}))

def test_from_callable_convert_to_beam_types(self):

def fn(

a: typing.List[int],

b: str = '',

*args: typing.Tuple[T_typing],

foo: typing.List[int],

**kwargs: typing.Dict[str, str]) -> typing.Tuple[typing.Any, ...]:

return a, b, args, foo, kwargs

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(

th.input_types,

((List[int], str, Tuple[T]), {

'foo': List[int], 'kwargs': Dict[str, str]

}))

self.assertEqual(th.output_types, ((Tuple[Any, ...], ), {}))

def test_from_callable_partial(self):

def fn(a: int) -> int:

return a

# functools.partial objects don't have __name__ attributes by default.

fn = functools.partial(fn, 1)

th = decorators.IOTypeHints.from_callable(fn)

self.assertRegex(th.debug_str(), r'unknown')

def test_from_callable_no_tagged_output(self):

def fn(x: int) -> str:

return str(x)

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(th.input_types, ((int, ), {}))

self.assertEqual(th.output_types, ((str, ), {}))

def fn2(x: int) -> typing.Iterable[str]:

yield str(x)

th = decorators.IOTypeHints.from_callable(fn2)

self.assertEqual(th.input_types, ((int, ), {}))

self.assertEqual(th.output_types, ((typehints.Iterable[str], ), {}))

def test_from_callable_tagged_output_union(self):

"""Tagged types are NOT extracted in from_callable. They stay embedded

in the main type and are extracted later in strip_iterable()."""

def fn(

x: int

) -> int | str | TaggedOutput[typing.Literal['errors'], float

| str] | TaggedOutput[

typing.Literal['warnings'], str]:

return x

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(th.input_types, ((int, ), {}))

# TaggedOutput members are preserved in the union no extraction yet.

output_type = th.output_types[0][0]

self.assertIsInstance(output_type, typehints.UnionConstraint)

self.assertEqual(th.output_types[1], {})

def test_from_callable_tagged_output_iterable(self):

"""Tagged types inside Iterable are preserved until strip_iterable."""

def fn(

x: int

) -> typing.Iterable[int | TaggedOutput[typing.Literal['errors'], str]]:

yield x

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(th.input_types, ((int, ), {}))

# The full Iterable[Union[int, TaggedOutput[...]]] is preserved.

output_type = th.output_types[0][0]

self.assertIsInstance(output_type, typehints.IterableTypeConstraint)

self.assertEqual(th.output_types[1], {})

def test_from_callable_tagged_output_only(self):

"""A standalone TaggedOutput annotation passes through from_callable."""

def fn(x: int) -> TaggedOutput[typing.Literal['errors'], str]:

pass

th = decorators.IOTypeHints.from_callable(fn)

self.assertEqual(th.input_types, ((int, ), {}))

# TaggedOutput[...] passes through convert_to_beam_type unchanged.

self.assertIs(typing.get_origin(th.output_types[0][0]), TaggedOutput)

self.assertEqual(th.output_types[1], {})

def test_getcallargs_forhints(self):

def fn(

a: int,

b: str = '',

*args: Tuple[T],

foo: List[int],

**kwargs: Dict[str, str]) -> Tuple[Any, ...]:

return a, b, args, foo, kwargs

callargs = decorators.getcallargs_forhints(fn, float, foo=List[str])

self.assertDictEqual(

callargs,

{

'a': float,

'b': str,

'args': Tuple[T],

'foo': List[str],

'kwargs': Dict[str, str]

})

def test_getcallargs_forhints_default_arg(self):

# Default args are not necessarily types, so they should be ignored.

def fn(a=List[int], b=None, *args, foo=(), **kwargs) -> Tuple[Any, ...]:

return a, b, args, foo, kwargs

callargs = decorators.getcallargs_forhints(fn)

self.assertDictEqual(

callargs,

{

'a': Any,

'b': Any,

'args': Tuple[Any, ...],

'foo': Any,

'kwargs': Dict[Any, Any]

})

def test_getcallargs_forhints_missing_arg(self):

def fn(a, b=None, *args, foo, **kwargs):

return a, b, args, foo, kwargs

with self.assertRaisesRegex(decorators.TypeCheckError, "missing.*'a'"):

decorators.getcallargs_forhints(fn, foo=List[int])

with self.assertRaisesRegex(decorators.TypeCheckError, "missing.*'foo'"):

decorators.getcallargs_forhints(fn, 5)

def test_origin_annotated(self):

def annotated(e: str) -> str:

return e

t = Map(annotated)

th = t.get_type_hints()

th = th.with_input_types(str)

self.assertRegex(th.debug_str(), r'with_input_types')

th = th.with_output_types(str)

self.assertRegex(

th.debug_str(),

r'(?s)with_output_types.*with_input_types.*Map.annotated')

class WithTypeHintsTest(unittest.TestCase):

def test_get_type_hints_no_settings(self):

class Base(WithTypeHints):

pass

th = Base().get_type_hints()

self.assertEqual(th.input_types, None)

self.assertEqual(th.output_types, None)

def test_get_type_hints_class_decorators(self):

@decorators.with_input_types(int, str)

@decorators.with_output_types(int)

class Base(WithTypeHints):

pass

th = Base().get_type_hints()

self.assertEqual(th.input_types, ((int, str), {}))

self.assertEqual(th.output_types, ((int, ), {}))

def test_get_type_hints_class_defaults(self):

class Base(WithTypeHints):

def default_type_hints(self):

return decorators.IOTypeHints(

input_types=((int, str), {}), output_types=((int, ), {}), origin=[])

th = Base().get_type_hints()

self.assertEqual(th.input_types, ((int, str), {}))

self.assertEqual(th.output_types, ((int, ), {}))

def test_get_type_hints_precedence_defaults_over_decorators(self):

@decorators.with_input_types(int)

@decorators.with_output_types(str)

class Base(WithTypeHints):

def default_type_hints(self):

return decorators.IOTypeHints(

input_types=((float, ), {}), output_types=None, origin=[])

th = Base().get_type_hints()

self.assertEqual(th.input_types, ((float, ), {}))

self.assertEqual(th.output_types, ((str, ), {}))

def test_get_type_hints_precedence_instance_over_defaults(self):

class Base(WithTypeHints):

def default_type_hints(self):

return decorators.IOTypeHints(

input_types=((float, ), {}), output_types=((str, ), {}), origin=[])

th = Base().with_input_types(int).get_type_hints()

self.assertEqual(th.input_types, ((int, ), {}))

self.assertEqual(th.output_types, ((str, ), {}))

def test_inherits_does_not_modify(self):

# See BEAM-8629.

@decorators.with_output_types(int)

class Subclass(WithTypeHints):

def __init__(self):

pass # intentionally avoiding super call

# These should be equal, but not the same object lest mutating the instance

# mutates the class.

self.assertIsNot(

Subclass()._get_or_create_type_hints(), Subclass._type_hints)

self.assertEqual(Subclass().get_type_hints(), Subclass._type_hints)

self.assertNotEqual(

Subclass().with_input_types(str)._type_hints, Subclass._type_hints)

class DecoratorsTest(unittest.TestCase):

def tearDown(self):

decorators._disable_from_callable = False

def test_disable_type_annotations(self):

self.assertFalse(decorators._disable_from_callable)

decorators.disable_type_annotations()

self.assertTrue(decorators._disable_from_callable)

def test_no_annotations_on_same_function(self):

def fn(a: int) -> int:

return a

with self.assertRaisesRegex(TypeCheckError,

r'requires .*int.* but was applied .*str'):

_ = ['a', 'b', 'c'] | Map(fn)

# Same pipeline doesn't raise without annotations on fn.

fn = decorators.no_annotations(fn)

_ = ['a', 'b', 'c'] | Map(fn)

def test_no_annotations_on_diff_function(self):

def fn(a: int) -> int:

return a

_ = [1, 2, 3] | Map(fn) # Doesn't raise - correct types.

with self.assertRaisesRegex(TypeCheckError,

r'requires .*int.* but was applied .*str'):

_ = ['a', 'b', 'c'] | Map(fn)

@decorators.no_annotations

def fn2(a: int) -> int:

return a

_ = ['a', 'b', 'c'] | Map(fn2) # Doesn't raise - no input type hints.

class ExtractTaggedFromTypeTest(unittest.TestCase):

"""Tests for _extract_tagged_from_type (Beam-level type extraction)."""

def test_simple_type_no_extraction(self):

main, tagged = decorators._extract_tagged_from_type(int)

self.assertEqual(main, int)

self.assertEqual(tagged, {})

def test_beam_union_no_tagged(self):

t = typehints.Union[int, str]

main, tagged = decorators._extract_tagged_from_type(t)

self.assertEqual(main, t)

self.assertEqual(tagged, {})

def test_standalone_tagged_output(self):

t = TaggedOutput[typing.Literal['errors'], str]

main, tagged = decorators._extract_tagged_from_type(t)

self.assertIs(main, decorators._NO_MAIN_TYPE)

self.assertEqual(tagged, {'errors': str})

def test_beam_union_with_tagged(self):

t = convert_to_beam_type(int | TaggedOutput[typing.Literal['errors'], str])

main, tagged = decorators._extract_tagged_from_type(t)

self.assertEqual(main, int)

self.assertEqual(tagged, {'errors': str})

def test_beam_union_multiple_tagged(self):

t = convert_to_beam_type(

int | TaggedOutput[typing.Literal['errors'], str]

| TaggedOutput[typing.Literal['warnings'], str])

main, tagged = decorators._extract_tagged_from_type(t)

self.assertEqual(main, int)

self.assertEqual(tagged, {'errors': str, 'warnings': str})

def test_beam_union_multiple_main_types(self):

t = convert_to_beam_type(

int | str | TaggedOutput[typing.Literal['errors'], bytes])

main, tagged = decorators._extract_tagged_from_type(t)

self.assertIsInstance(main, typehints.UnionConstraint)

self.assertIn(int, main.union_types)

self.assertIn(str, main.union_types)

self.assertEqual(tagged, {'errors': bytes})

def test_beam_union_tagged_only(self):

t = convert_to_beam_type(

TaggedOutput[typing.Literal['errors'], str]

| TaggedOutput[typing.Literal['warnings'], int])

main, tagged = decorators._extract_tagged_from_type(t)

self.assertIs(main, decorators._NO_MAIN_TYPE)

self.assertEqual(tagged, {'errors': str, 'warnings': int})

def test_bare_tagged_output_standalone(self):

with self.assertLogs(level='WARNING') as cm:

main, tagged = decorators._extract_tagged_from_type(TaggedOutput)

self.assertIn('Bare TaggedOutput will be ignored', cm.output[0])

self.assertIs(main, decorators._NO_MAIN_TYPE)

self.assertEqual(tagged, {})

def test_bare_tagged_output_in_union(self):

with self.assertLogs(level='WARNING') as cm:

t = convert_to_beam_type(str | TaggedOutput)

main, tagged = decorators._extract_tagged_from_type(t)

self.assertIn('Bare TaggedOutput will be ignored', cm.output[0])

self.assertEqual(main, str)

self.assertEqual(tagged, {})

if __name__ == '__main__':

unittest.main()