Files
pytorch/test/jit/test_decorator.py
Anthony Barbier bf7e290854 Add __main__ guards to jit tests (#154725)
This PR is part of a series attempting to re-submit https://github.com/pytorch/pytorch/pull/134592 as smaller PRs.

In jit tests:

- Add and use a common raise_on_run_directly method for when a user runs a test file directly which should not be run this way. Print the file which the user should have run.
- Raise a RuntimeError on tests which have been disabled (not run)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/154725
Approved by: https://github.com/clee2000
2025-06-16 10:28:45 +00:00

29 lines
834 B
Python

# Owner(s): ["oncall: jit"]
import torch
from jit.myfunction_a import my_function_a
from torch.testing._internal.jit_utils import JitTestCase
class TestDecorator(JitTestCase):
def test_decorator(self):
# Note: JitTestCase.checkScript() does not work with decorators
# self.checkScript(my_function_a, (1.0,))
# Error:
# RuntimeError: expected def but found '@' here:
# @my_decorator
# ~ <--- HERE
# def my_function_a(x: float) -> float:
# Do a simple torch.jit.script() test instead
fn = my_function_a
fx = torch.jit.script(fn)
self.assertEqual(fn(1.0), fx(1.0))
if __name__ == "__main__":
raise RuntimeError(
"This test is not currently used and should be "
"enabled in discover_tests.py if required."
)