# mypy: ignore-errors

# Owner(s): ["module: unknown"]

import token

from itertools import product

from pathlib import Path

import torch

from torch.testing._internal.common_utils import (

instantiate_parametrized_tests,

parametrize,

run_tests,

TestCase,

)

MM = "@"

BINARY_RETURNS_BOOL = "!=", "<", "<=", "==", ">", ">="

BINARY_ACCEPTS_FLOAT_OR_INT = "%", "*", "**", "+", "-", "/", "//"

BINARY_ACCEPTS_INT_ONLY = "&", "<<", ">>", "^", "|"

BINARY_OPS = (

*BINARY_RETURNS_BOOL,

*BINARY_ACCEPTS_FLOAT_OR_INT,

*BINARY_ACCEPTS_INT_ONLY,

MM,

)

BINARY_RETURNS_FLOAT = ("/",)

UNARY_ACCEPTS_FLOAT_OR_INT = "+", "-"

UNARY_ACCEPTS_INT_ONLY = ("~",)

UNARY_OPS = *UNARY_ACCEPTS_FLOAT_OR_INT, *UNARY_ACCEPTS_INT_ONLY

PUNCTUATION = ",", ";"

OPERATORS = *UNARY_OPS, *BINARY_OPS, *PUNCTUATION

FLOATS = 1.5, torch.tensor((2.5, 3.5))

INTS = 3, torch.tensor((1, 2))

ALL = *FLOATS, *INTS

TYPE_TEST_FILE = Path(__file__).parent / "pass/arithmetic_ops.py"

class TestPythonOperators(TestCase):

# Prove that UNARY_OPS, BINARY_OPS, and OPERATORS are correct and complete

def test_operators_are_correct_and_complete(self):

self.assertFalse(set(OPERATORS).difference(token.EXACT_TOKEN_TYPES))

unary, binary, punctuation = {}, {}, {}

for op in token.EXACT_TOKEN_TYPES:

if op in PUNCTUATION:

punctuation[op] = True

else:

try:

unary[op] = compile(f"{op}1 ; {op}a", op, "single")

except SyntaxError:

pass

try:

binary[op] = compile(f"2 {op} 3 ; a {op} b", op, "single")

except SyntaxError:

pass

self.assertEqual(sorted(unary), sorted(UNARY_OPS))

self.assertEqual(sorted(binary), sorted(BINARY_OPS))

self.assertEqual(sorted(punctuation), sorted(PUNCTUATION))

def test_type_tests_are_complete(self):

binary, unary = {}, []

with TYPE_TEST_FILE.open() as fp:

# Looking for lines like: assert_type(TENSOR ^ BOOL, Tensor)

# But not: assert_type(BOOL ^ BINARY, Binary)

lines = (i for i in fp if "TENSOR" in i)

for line in lines:

if expr := line.partition("assert_type(")[2].partition(",")[0]:

if expr[0].isalpha():

# ** formats differently from all other operators

a, op, b = expr.replace("**", " ** ").split()

binary.setdefault(op, []).append((a, b))

else:

unary.append(expr[0])

self.assertEqual(sorted(unary), sorted(UNARY_OPS))

self.assertEqual(sorted(binary), sorted(BINARY_OPS))

value, *values = binary.values()

self.assertEqual(values, [value] * len(values))

@parametrize("a, op, b", product(ALL, BINARY_OPS, ALL))

def test_binary(self, a, op, b):

try:

r = eval(f"a {op} b")

except Exception as e:

r = e

any_tensor = isinstance(a, torch.Tensor) or isinstance(b, torch.Tensor)

any_float = _any_float(a, b)

returns_float = any_float or op in BINARY_RETURNS_FLOAT

if op == MM:

if not (isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor)):

self.assertIsInstance(r, TypeError)

elif a is b:

self.assertIsInstance(r, torch.Tensor)

else:

self.assertIsInstance(r, RuntimeError)

elif any_tensor:

if op in BINARY_ACCEPTS_INT_ONLY and any_float:

# See https://github.com/pytorch/pytorch/issues/15754

self.assertIsInstance(r, NotImplementedError)

else:

self.assertIsInstance(r, torch.Tensor)

if op in BINARY_RETURNS_BOOL:

self.assertEqual(r.dtype, torch.bool)

elif op in BINARY_ACCEPTS_INT_ONLY:

self.assertFalse(r.dtype.is_floating_point)

elif op in BINARY_ACCEPTS_FLOAT_OR_INT:

self.assertEqual(r.dtype.is_floating_point, returns_float)

else:

self.assertFalse("Logic error")

elif op in BINARY_RETURNS_BOOL:

self.assertIsInstance(r, bool)

elif op in BINARY_ACCEPTS_INT_ONLY:

if any_float:

self.assertIsInstance(r, TypeError)

else:

self.assertIsInstance(r, int)

elif returns_float:

self.assertIsInstance(r, float)

else:

self.assertIsInstance(r, int)

@parametrize("op, a", product(UNARY_OPS, ALL))

def test_unary(self, op, a):

try:

r = eval(f"{op} a")

except Exception as e:

r = e

if op in UNARY_ACCEPTS_INT_ONLY and _any_float(a):

self.assertIsInstance(r, TypeError)

elif isinstance(a, torch.Tensor):

self.assertIsInstance(r, torch.Tensor)

elif op in UNARY_ACCEPTS_INT_ONLY:

self.assertIsInstance(r, int)

elif isinstance(a, float):

self.assertIsInstance(r, float)

else:

self.assertIsInstance(r, int)

def _any_float(*x):

for i in x:

if isinstance(i, float) or (

isinstance(i, torch.Tensor) and i.dtype.is_floating_point

):

return True

return False

instantiate_parametrized_tests(TestPythonOperators)

if __name__ == "__main__":

run_tests()