github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/python_callable_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  
    17  import os
    18  import unittest
    19  
    20  import apache_beam as beam
    21  from apache_beam.utils.python_callable import PythonCallableWithSource
    22  
    23  
    24  class PythonCallableWithSourceTest(unittest.TestCase):
    25    def test_builtin(self):
    26      self.assertEqual(PythonCallableWithSource.load_from_source('str'), str)
    27  
    28    def test_builtin_attribute(self):
    29      self.assertEqual(
    30          PythonCallableWithSource.load_from_source('str.lower'), str.lower)
    31  
    32    def test_fully_qualified_name(self):
    33      self.assertEqual(
    34          PythonCallableWithSource.load_from_source('os.path.abspath'),
    35          os.path.abspath)
    36  
    37    def test_expression(self):
    38      self.assertEqual(PythonCallableWithSource('lambda x: x*x')(10), 100)
    39  
    40    def test_expression_with_dependency(self):
    41      self.assertEqual(
    42          PythonCallableWithSource('import math\nlambda x: math.sqrt(x) + x')(
    43              100),
    44          110)
    45  
    46    def test_def(self):
    47      self.assertEqual(
    48          PythonCallableWithSource(
    49              """
    50              def foo(x):
    51                  return x * x
    52          """)(10),
    53          100)
    54  
    55    def test_def_with_preamble(self):
    56      self.assertEqual(
    57          PythonCallableWithSource(
    58              """
    59              def bar(x):
    60                  return x + 1
    61  
    62              def foo(x):
    63                  return bar(x) * x
    64          """)(10),
    65          110)
    66  
    67    def test_class(self):
    68      self.assertEqual(
    69          PythonCallableWithSource(
    70              """
    71              class BareClass:
    72                def __init__(self, x):
    73                  self.x = x
    74          """)(10).x,
    75          10)
    76  
    77      self.assertEqual(
    78          PythonCallableWithSource(
    79              """
    80              class SubClass(object):
    81                def __init__(self, x):
    82                  self.x = x
    83          """)(10).x,
    84          10)
    85  
    86    def test_pycallable_map(self):
    87      p = beam.Pipeline()
    88      result = (
    89          p
    90          | beam.Create([1, 2, 3])
    91          | beam.Map(PythonCallableWithSource("lambda x: x")))
    92      self.assertEqual(result.element_type, int)
    93  
    94  
    95  if __name__ == '__main__':
    96    unittest.main()