Fix distinct error codes test (#73)

Closes #71
pytorch-labs · Sep 5, 2024 · c755731 · c755731
1 parent 555bb8d
commit c755731
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 20 deletions.
diff --git a/tests/test_torchfix.py b/tests/test_torchfix.py
@@ -65,8 +65,8 @@ def test_errorcodes_distinct():
         LOGGER.info("Checking error code for %s", visitor.__class__.__name__)
         errors = visitor.ERRORS
         for e in errors:
-            assert e not in seen
-            seen.add(e)
+            assert e.error_code not in seen
+            seen.add(e.error_code)
 
 
 def test_parse_error_code_str():

diff --git a/torchfix/visitors/vision/singleton_import.py b/torchfix/visitors/vision/singleton_import.py
@@ -9,22 +9,8 @@ class TorchVisionSingletonImportVisitor(TorchVisitor):
         TorchError(
             "TOR203",
             (
-                "Consider replacing 'import torchvision.datasets as datasets' "
-                "with 'from torchvision import datasets'."
-            ),
-        ),
-        TorchError(
-            "TOR203",
-            (
-                "Consider replacing 'import torchvision.models as models' "
-                "with 'from torchvision import models'."
-            ),
-        ),
-        TorchError(
-            "TOR203",
-            (
-                "Consider replacing 'import torchvision.transforms as transforms' "
-                "with 'from torchvision import transforms'."
+                "Consider replacing 'import torchvision.{module} as {module}' "
+                "with 'from torchvision import {module}'."
             ),
         ),
     ]
@@ -53,8 +39,8 @@ def visit_Import(self, node: cst.Import) -> None:
                         )
                     self.add_violation(
                         node,
-                        error_code=self.ERRORS[i].error_code,
-                        message=self.ERRORS[i].message(),
+                        error_code=self.ERRORS[0].error_code,
+                        message=self.ERRORS[0].message(module=import_attr),
                         replacement=replacement,
                     )
                 break