Files
nvidia--tensorrt/tools/Polygraphy/tests/func/test_func.py
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

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()