Merge pull request #5 from pylint-dev/add-tensorflow-import-checker

Add tensorflow import checker

Author: PeterHamfelt

.gitignore

@@ -157,12 +157,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-.idea/
-
-# Mr Developer
-.mr.developer.cfg
-.project
-.pydevproject
-.idea
-env.txt
-poetry.lock
+# .idea/

CONTRIBUTORS.txt

@@ -15,14 +15,14 @@
 class PandasImportChecker(BaseChecker):
     name = "pandas-import"
     msgs = {
-        "W8001": (
+        "W8101": (
             "Pandas imported with incorrect alias",
             "pandas-import",
             "Pandas should be imported with the alias `pd` to maintain consistency with common practices. "
             "Importing pandas with any other alias can lead to confusion. "
             "Consider using `import pandas as pd` for clarity and adherence to the convention.",
         ),
-        "W8002": (
+        "W8102": (
             "Direct import from Pandas discouraged",
             "pandas-importfrom",
             "Direct imports from Pandas using `from pandas import ...` are discouraged to maintain code clarity and "

pylint_ml/checkers/pandas/import_pandas.py

@@ -0,0 +1,44 @@
+# 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 tensorflow 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 TensorflowImportChecker(BaseChecker):
+    name = "tensorflow-import"
+    msgs = {
+        "W8401": (
+            "Tensorflow imported with incorrect alias",
+            "tensorflow-import",
+            "Tensorflow should be imported with the alias `tf` to maintain consistency with common practices. "
+            "Importing Tensorflow with any other alias can lead to confusion. "
+            "Consider using `import tensorflow as tf` for clarity and adherence to the convention.",
+        ),
+        "W8402": (
+            "Direct import from Tensorflow discouraged",
+            "tensorflow-importfrom",
+            "Direct imports from Tensorflow using `from tensorflow import ...` are discouraged to maintain code "
+            "clarity and prevent potential conflicts. Tensorflow should be imported with the alias `tf` following "
+            "common practices. Using any other alias or direct import method can lead to confusion. "
+            "Consider using `import tensorflow as tf` to adhere to the convention and ensure consistency.",
+        ),
+    }
+
+    @only_required_for_messages("tensorflow-import")
+    def visit_import(self, node: nodes.Import) -> None:
+        for name, alias in node.names:
+            if name == "tensorflow" and alias != "tf":
+                self.add_message("tensorflow-import", node=node, confidence=HIGH)
+
+    @only_required_for_messages("tensorflow-importfrom")
+    def visit_importfrom(self, node: nodes.ImportFrom) -> None:
+        if node.modname == "tensorflow":
+            self.add_message("tensorflow-importfrom", node=node, confidence=HIGH)

pylint_ml/checkers/tensorflow/import_tensorflow.py

@@ -0,0 +1,53 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.tensorflow.import_tensorflow import TensorflowImportChecker
+
+
+class TestTensorflowImport(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = TensorflowImportChecker
+
+    def test_correct_tensorflow_import(self):
+        tensorflow_import_node = astroid.extract_node(
+            """
+        import tensorflow as tf
+        """
+        )
+
+        with self.assertNoMessages():
+            self.checker.visit_import(tensorflow_import_node)
+
+    def test_incorrect_tensorflow_import(self):
+        tensorflow_import_node = astroid.extract_node(
+            """
+        import tensorflow as tflow
+        """
+        )
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="tensorflow-import",
+                confidence=HIGH,
+                node=tensorflow_import_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_import(tensorflow_import_node)
+
+    def test_incorrect_tensorflow_import_from(self):
+        tensorflow_importfrom_node = astroid.extract_node(
+            """
+        from tensorflow import math
+        """
+        )
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="tensorflow-importfrom",
+                confidence=HIGH,
+                node=tensorflow_importfrom_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_importfrom(tensorflow_importfrom_node)