support selection of tests that use op infos (#17)

Author: pmeier

pytest_pytorch/plugin.py

@@ -1,13 +1,12 @@
+import inspect
 import re
 import unittest.mock
 import warnings
-from typing import Pattern
 
 from _pytest.unittest import TestCaseFunction, UnitTestCase
 
 try:
-    from torch.testing._internal.common_device_type import get_device_type_test_bases
-    from torch.testing._internal.common_utils import TestCase as PyTorchTestCaseTemplate
+    from torch.testing._internal.common_utils import TestCase as TestCaseTemplate
 
     TORCH_AVAILABLE = True
 except ImportError:
@@ -17,93 +16,73 @@
         "Disabling the plugin 'pytest-pytorch', because 'torch' could not be imported."
     )
 
-    def get_device_type_test_bases():
-        return []
 
-    class PyTorchTestCaseTemplate:
-        pass
+class TemplatedName(str):
+    def __new__(cls, name, template_name):
+        self = super().__new__(cls, name)
+        self._template_name = template_name
+        return self
 
+    def __eq__(self, other):
+        exact_match = str.__eq__(self, other)
+        if exact_match:
+            return True
 
-class PytestPyTorchInternalError(Exception):
-    def __init__(self, msg):
-        super().__init__(
-            f"{msg}\n"
-            f"This is an internal error of the pytest plugin 'pytest-pytorch'."
-            f"If you encounter this during normal operation, please file an issue "
-            f"https://github.com/Quansight/pytest-pytorch/issues."
-        )
-
-
-TEST_BASE_DEVICE_PATTERN = re.compile(r"(?P<device>\w*?)TestBase$")
-
-
-def _get_devices():
-    devices = []
-    for test_base in get_device_type_test_bases():
-        match = TEST_BASE_DEVICE_PATTERN.match(test_base.__name__)
-        if not match:
-            raise PytestPyTorchInternalError(
-                f"Unable to extract device name from {test_base.__name__}"
-            )
-
-        devices.append(match.group("device"))
+        if not self._template_name:
+            return False
 
-    return devices
+        return str.__eq__(self._template_name, other)
 
+    def __hash__(self):
+        return super().__hash__()
 
-DEVICES = _get_devices()
 
+class TemplatedTestCaseFunction(TestCaseFunction):
+    _TEMPLATE_NAME_PATTERN = re.compile(r"def (?P<template_name>test_\w+)\(")
 
-class TemplateName(str):
-    _TEMPLATE_NAME_PATTERN: Pattern
+    @classmethod
+    def _extract_template_name(cls, callobj):
+        if not callobj:
+            return None
 
-    def __init__(self, _):
-        super().__init__()
-        match = self._TEMPLATE_NAME_PATTERN.match(self)
+        match = cls._TEMPLATE_NAME_PATTERN.search(inspect.getsource(callobj))
         if not match:
-            raise PytestPyTorchInternalError(
-                f"Unable to extract template name from {self}"
-            )
-        self._template_name = match.group("template_name")
-
-    def __eq__(self, other):
-        return str.__eq__(self, other) or str.__eq__(self._template_name, other)
-
-    def __hash__(self) -> int:
-        return super().__hash__()
-
-
-class TestCaseFunctionTemplateName(TemplateName):
-    _TEMPLATE_NAME_PATTERN = re.compile(
-        fr"(?P<template_name>\w*?)_({'|'.join([device.lower() for device in DEVICES])})"
-    )
+            return None
 
+        return match.group("template_name")
 
-class PyTorchTestCaseFunction(TestCaseFunction):
     @classmethod
-    def from_parent(cls, parent, *, name, **kw):
+    def from_parent(cls, parent, *, name, callobj, **kw):
         return super().from_parent(
-            parent, name=TestCaseFunctionTemplateName(name), **kw
+            parent, name=TemplatedName(name, cls._extract_template_name(callobj)), **kw
         )
 
 
-class TestCaseTemplateName(TemplateName):
-    _TEMPLATE_NAME_PATTERN = re.compile(
-        fr"(?P<template_name>\w*?)({'|'.join([device.upper() for device in DEVICES])})"
-    )
+class TemplatedTestCase(UnitTestCase):
+    @classmethod
+    def _extract_template_name(cls, name, obj):
+        if not obj:
+            return None
+
+        if not hasattr(obj, "device_type"):
+            return None
 
+        return name[: -len(obj.device_type)]
 
-class PyTorchTestCase(UnitTestCase):
     @classmethod
     def from_parent(cls, parent, *, name, obj=None):
-        return super().from_parent(parent, name=TestCaseTemplateName(name), obj=obj)
+        return super().from_parent(
+            parent,
+            name=TemplatedName(name, cls._extract_template_name(name, obj)),
+            obj=obj,
+        )
 
     def collect(self):
         # Yes, this is a bad practice. Unfortunately, there is no other option to
         # inject our custom 'TestCaseFunction' without duplicating everything in
         # 'UnitTestCase.collect()'
         with unittest.mock.patch(
-            "_pytest.unittest.TestCaseFunction", new=PyTorchTestCaseFunction
+            "_pytest.unittest.TestCaseFunction", new=TemplatedTestCaseFunction
         ):
             yield from super().collect()
 
@@ -113,12 +92,9 @@ def pytest_pycollect_makeitem(collector, name, obj):
         return None
 
     try:
-        if (
-            not issubclass(obj, PyTorchTestCaseTemplate)
-            or obj is PyTorchTestCaseTemplate
-        ):
+        if not issubclass(obj, TestCaseTemplate) or obj is TestCaseTemplate:
             return None
     except Exception:
         return None
 
-    return PyTorchTestCase.from_parent(collector, name=name, obj=obj)
+    return TemplatedTestCase.from_parent(collector, name=name, obj=obj)

tests/assets/test_op_infos.py

@@ -0,0 +1,18 @@
+import torch
+from torch.testing._internal.common_device_type import (
+    dtypes,
+    instantiate_device_type_tests,
+    ops,
+)
+from torch.testing._internal.common_methods_invocations import OpInfo
+from torch.testing._internal.common_utils import TestCase
+
+
+class TestFoo(TestCase):
+    @dtypes(torch.float16, torch.int32)
+    @ops([OpInfo("add"), OpInfo("sub")])
+    def test_bar(self, device, dtype, op):
+        assert True
+
+
+instantiate_device_type_tests(TestFoo, globals())

tests/test_plugin.py

@@ -205,3 +205,30 @@ def test_nested_names(testdir, file, cmds, outcomes):
     testdir.copy_example(file)
     result = testdir.runpytest(*cmds)
     result.assert_outcomes(**outcomes)
+
+
+@make_params(
+    "test_op_infos.py",
+    Config(
+        "*testcase-*test-*op-*device-*dtype",
+        new_cmds=(),
+        legacy_cmds=(),
+        passed=8,
+    ),
+    Config(
+        "1testcase-*test-*op-*device-*dtype",
+        new_cmds="::TestFoo",
+        legacy_cmds=("-k", "TestFoo"),
+        passed=8,
+    ),
+    Config(
+        "1testcase-1test-*op-*device-*dtype",
+        new_cmds="::TestFoo::test_bar",
+        legacy_cmds=("-k", "TestFoo and test_bar"),
+        passed=8,
+    ),
+)
+def test_op_infos(testdir, file, cmds, outcomes):
+    testdir.copy_example(file)
+    result = testdir.runpytest(*cmds)
+    result.assert_outcomes(**outcomes)

tox.ini

@@ -31,6 +31,8 @@ deps =
   torch
   # The nightlies do not specify numpy as requirement
   numpy
+  # Importing OpInfo requires scipy unconditionally
+  scipy
 commands =
   pytest -c pytest.ini {posargs}