Merge pull request #7 from pylint-dev/6-add-torch-import-checker

Add torch import checker

Author: PeterHamfelt

pylint_ml/checkers/torch/__init__.py

@@ -0,0 +1,43 @@
+# Licensed under the MIT: https://mit-license.org/
+# For details: https://github.com/pylint-dev/pylint-ml/LICENSE
+# Copyright (c) https://github.com/pylint-dev/pylint-ml/CONTRIBUTORS.txt
+
+"""Check for import of torch library."""
+
+from __future__ import annotations
+
+from astroid import nodes
+from pylint.checkers import BaseChecker
+from pylint.checkers.utils import only_required_for_messages
+from pylint.interfaces import HIGH
+
+
+class TorchImportChecker(BaseChecker):
+    name = "torch-import"
+    msgs = {
+        "W8401": (
+            "Torch imported with alias",
+            "torch-import",
+            "Torch should be imported without an alias to maintain consistency with common practices. "
+            "Importing Torch with an alias can lead to confusion. "
+            "Consider using `import torch` for clarity and adherence to the convention.",
+        ),
+        "W8402": (
+            "Direct import from Torch discouraged",
+            "torch-importfrom",
+            "Direct imports from Torch using `from torch import ...` are discouraged to maintain code "
+            "clarity and prevent potential conflicts. Using any alias or direct import method can lead to confusion. "
+            "Consider using `import torch` to adhere to the convention and ensure consistency.",
+        ),
+    }
+
+    @only_required_for_messages("torch-import")
+    def visit_import(self, node: nodes.Import) -> None:
+        for name, alias in node.names:
+            if name == "torch" and alias:  # Alias is used
+                self.add_message("torch-import", node=node, confidence=HIGH)
+
+    @only_required_for_messages("torch-importfrom")
+    def visit_importfrom(self, node: nodes.ImportFrom) -> None:
+        if node.modname == "torch":
+            self.add_message("torch-importfrom", node=node, confidence=HIGH)

pylint_ml/checkers/torch/import_torch.py

@@ -4,6 +4,8 @@
 
 from pylint_ml.checkers.numpy.import_numpy import NumpyImportChecker
 from pylint_ml.checkers.pandas.import_pandas import PandasImportChecker
+from pylint_ml.checkers.tensorflow.import_tensorflow import TensorflowImportChecker
+from pylint_ml.checkers.torch.import_torch import TorchImportChecker
 
 
 def register(linter: PyLinter) -> None:
@@ -13,3 +15,12 @@ def register(linter: PyLinter) -> None:
 
     # Pandas
     linter.register_checker(PandasImportChecker(linter))
+
+    # Tensorflow
+    linter.register_checker(TensorflowImportChecker(linter))
+
+    # Torch
+    linter.register_checker(TorchImportChecker(linter))
+
+    # Sklearn
+    # Scipy

pylint_ml/plugin.py

@@ -0,0 +1,53 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.torch.import_torch import TorchImportChecker
+
+
+class TestTorchImport(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = TorchImportChecker
+
+    def test_correct_torch_import(self):
+        torch_import_node = astroid.extract_node(
+            """
+        import torch
+        """
+        )
+
+        with self.assertNoMessages():
+            self.checker.visit_import(torch_import_node)
+
+    def test_incorrect_torch_import(self):
+        torch_import_node = astroid.extract_node(
+            """
+        import torch as th
+        """
+        )
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="torch-import",
+                confidence=HIGH,
+                node=torch_import_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_import(torch_import_node)
+
+    def test_incorrect_torch_import_from(self):
+        torch_importfrom_node = astroid.extract_node(
+            """
+        from torch import min
+        """
+        )
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="torch-importfrom",
+                confidence=HIGH,
+                node=torch_importfrom_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_importfrom(torch_importfrom_node)