chore: import upstream snapshot with attribution

tools/Polygraphy/tests/func/test_func.py

@@ -0,0 +1,187 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed 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.
+#
+import pytest
+from polygraphy import func
+from polygraphy.exception import PolygraphyException, PolygraphyInternalException
+
+
+class TestExtend:
+    def test_override_rv(self):
+        def x():
+            return 1
+
+        # Since y explicitly returns something, the return value of x is discarded.
+        @func.extend(x)
+        def y(elem):
+            assert elem == 1
+            return 2
+
+        assert y() == 2
+
+    def test_extend_named_parameters(self):
+        def x(arg0, arg1):
+            return arg0, arg1
+
+        @func.extend(x)
+        def y(elem0, elem1):
+            pass
+
+        arg0, arg1 = y(arg1=1, arg0=0)
+        assert arg0 == 0
+        assert arg1 == 1
+
+    def test_extend_0_args_1_rv(self):
+        def x():
+            return 1
+
+        @func.extend(x)
+        def y(elem):
+            assert elem == 1
+
+        assert y() == 1
+
+    def test_extend_0_args_2_rv(self):
+        def x():
+            return 1, 2
+
+        @func.extend(x)
+        def y(elem0, elem1):
+            assert elem0 == 1
+            assert elem1 == 2
+
+        assert y() == (1, 2)
+
+    def test_extend_1_args_0_rv(self):
+        def x(arg0):
+            pass
+
+        @func.extend(x)
+        def y():
+            pass
+
+        y(1)
+
+    def test_extend_1_args_1_rv(self):
+        def x(arg0):
+            assert arg0 == 1
+            return 3
+
+        @func.extend(x)
+        def y(elem):
+            assert elem == 3
+
+        assert y(1) == 3
+
+    def test_extend_2_args_2_rv(self):
+        def x(arg0, arg1):
+            assert arg0 == -1
+            assert arg1 == -1
+            return 1, 2
+
+        @func.extend(x)
+        def y(elem0, elem1):
+            assert elem0 == 1
+            assert elem1 == 2
+
+        assert y(-1, -1) == (1, 2)
+
+    def test_extend_can_modify_rv(self):
+        def x():
+            return []
+
+        @func.extend(x)
+        def y(lst):
+            lst.extend([1, 2, 3])
+
+        assert x() == []
+        assert y() == [1, 2, 3]
+
+    def test_extend_can_modify_rv_objects(self):
+        class ModifiableObj:
+            def __init__(self):
+                self.value = 0
+
+        def x():
+            return ModifiableObj()
+
+        @func.extend(x)
+        def y(mo):
+            mo.value = 1
+
+        assert x().value == 0
+        assert y().value == 1
+
+    def test_extend_incorrect_num_args(self):
+        def x():
+            return 1, 2
+
+        with pytest.raises(
+            PolygraphyException,
+            match=r"Function: y accepts 1 parameter\(s\), but needs to accept 2 parameter\(s\)",
+        ):
+
+            @func.extend(x)
+            def y(elem0):
+                assert elem0 == 1
+
+            y()
+
+    @pytest.mark.parametrize("args_mode", ["kwargs", "args", "mixed"])
+    def test_extend_forward_parameters(self, args_mode):
+        def x(x_arg0, x_arg1):
+            return x_arg0 + x_arg1
+
+        @func.extend(x)
+        def y(x_arg0, x_arg1, x_ret):
+            assert x_ret == x_arg0 + x_arg1
+
+        if args_mode == "kwargs":
+            assert y(x_arg0=2, x_arg1=1) == 3
+        elif args_mode == "args":
+            assert y(2, 1) == 3
+        else:
+            assert args_mode == "mixed"
+            assert y(2, x_arg1=1) == 3
+
+
+class TestConstantMethod:
+    def test_cannot_modify_attrs(self):
+        class Dummy:
+            def __init__(self):
+                self.x = 1
+
+            @func.constantmethod
+            def modify_x(self):
+                self.x = 2
+
+        d = Dummy()
+        with pytest.raises(
+            PolygraphyInternalException, match="was mutated in a constant method"
+        ):
+            d.modify_x()
+
+    def test_cannot_add_attrs(self):
+        class Dummy:
+            @func.constantmethod
+            def modify_x(self):
+                self.x = 2
+
+        d = Dummy()
+        with pytest.raises(
+            PolygraphyInternalException, match="was mutated in a constant method"
+        ):
+            d.modify_x()