188 lines
4.5 KiB
Python
188 lines
4.5 KiB
Python
#
|
|
# 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()
|