Add sklearn import checker

Author: PeterHamfelt

pylint_ml/checkers/sklearn/import_sklearn.py

@@ -0,0 +1,32 @@
+# 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 sklearn 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 SklearnImportChecker(BaseChecker):
+    name = "sklearn-import"
+    msgs = {
+        "W8401": (
+            "Direct or aliased Sklearn import detected",
+            "sklearn-import",
+            "Using `import sklearn` or `import sklearn as ...` is not recommended. For better clarity and consistency, "
+            "it is advisable to import specific submodules directly, using the `from sklearn import ...` syntax. "
+            "This approach prevents confusion and aligns with common practices by explicitly stating which "
+            "components of Sklearn are being used.",
+        ),
+    }
+
+    @only_required_for_messages("sklearn-import")
+    def visit_import(self, node: nodes.Import) -> None:
+        for name, _ in node.names:
+            if name == "sklearn":
+                self.add_message("sklearn-import", node=node, confidence=HIGH)

pylint_ml/plugin.py

@@ -5,6 +5,7 @@
 from pylint_ml.checkers.numpy.import_numpy import NumpyImportChecker
 from pylint_ml.checkers.pandas.import_pandas import PandasImportChecker
 from pylint_ml.checkers.scipy.import_scipy import ScipyImportChecker
+from pylint_ml.checkers.sklearn.import_sklearn import SklearnImportChecker
 from pylint_ml.checkers.tensorflow.import_tensorflow import TensorflowImportChecker
 from pylint_ml.checkers.torch.import_torch import TorchImportChecker
 
@@ -27,6 +28,7 @@ def register(linter: PyLinter) -> None:
     linter.register_checker(ScipyImportChecker(linter))
 
     # Sklearn
+    linter.register_checker(SklearnImportChecker(linter))
 
     # Theano
     # Matplotlib

tests/checkers/test_sklearn/__init__.py

@@ -0,0 +1,36 @@
+import astroid
+import pylint.testutils
+from pylint.interfaces import HIGH
+
+from pylint_ml.checkers.sklearn.import_sklearn import SklearnImportChecker
+
+
+class TestSklearnImport(pylint.testutils.CheckerTestCase):
+    CHECKER_CLASS = SklearnImportChecker
+
+    def test_correct_sklearn_import(self):
+        sklearn_import_node = astroid.extract_node(
+            """
+        from sklearn import datasets
+        """
+        )
+
+        with self.assertNoMessages():
+            self.checker.visit_import(sklearn_import_node)
+
+    def test_incorrect_sklearn_import(self):
+        sklearn_import_node = astroid.extract_node(
+            """
+        import sklearn as skl
+        """
+        )
+
+        with self.assertAddsMessages(
+            pylint.testutils.MessageTest(
+                msg_id="sklearn-import",
+                confidence=HIGH,
+                node=sklearn_import_node,
+            ),
+            ignore_position=True,
+        ):
+            self.checker.visit_import(sklearn_import_node)