Add sklearn contrib (#44)

Author: cthoyt

docs/source/conf.py

@@ -232,6 +232,7 @@
     'torch': ('https://pytorch.org/docs/stable', None),
     'numpy': ('https://numpy.org/doc/stable', None),
     'optuna': ('https://optuna.readthedocs.io/en/latest', None),
+    'sklearn': ('https://scikit-learn.org/stable', None),
 }
 
 autoclass_content = 'both'

docs/source/contrib/index.rst

@@ -9,3 +9,4 @@ External Package Integrations
    torch
    numpy
    optuna
+   sklearn

docs/source/contrib/sklearn.rst

@@ -0,0 +1,4 @@
+Scikit-learn
+============
+.. automodule:: class_resolver.contrib.sklearn
+    :members:

setup.cfg

@@ -82,6 +82,8 @@ optuna =
     optuna
 numpy =
     numpy
+sklearn =
+    scikit-learn
 
 [options.entry_points]
 class_resolver_demo =

src/class_resolver/contrib/sklearn.py

@@ -0,0 +1,73 @@
+"""
+Scikit-learn is a generic machine learning package with implementations of
+algorithms for classification, regression, dimensionality reduction, clustering,
+as well as other generic tooling.
+
+The ``class-resolver`` provides several class resolvers for instantiating various
+implementations, such as those of linear models.
+"""  # noqa:D205,D400
+
+from sklearn.base import BaseEstimator
+from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
+from sklearn.linear_model import (
+    LogisticRegression,
+    LogisticRegressionCV,
+    PassiveAggressiveClassifier,
+    Perceptron,
+    RidgeClassifier,
+    RidgeClassifierCV,
+    SGDClassifier,
+)
+from sklearn.tree import DecisionTreeClassifier
+
+from ..api import ClassResolver
+
+__all__ = [
+    "classifier_resolver",
+]
+
+classifier_resolver: ClassResolver[BaseEstimator] = ClassResolver(
+    [
+        LogisticRegression,
+        LogisticRegressionCV,
+        PassiveAggressiveClassifier,
+        Perceptron,
+        RidgeClassifier,
+        RidgeClassifierCV,
+        SGDClassifier,
+        DecisionTreeClassifier,
+        RandomForestClassifier,
+        GradientBoostingClassifier,
+    ],
+    base=BaseEstimator,
+    base_as_suffix=False,
+    default=LogisticRegression,
+)
+"""A resolver for classifiers.
+
+The default value is :class:`sklearn.linear_model.LogisticRegression`.
+This resolver can be used like in the following:
+
+.. code-block:: python
+
+    from sklearn import datasets
+    from sklearn.model_selection import train_test_split
+
+    from class_resolver.contrib.sklearn import classifier_resolver
+
+    # Prepare a dataset
+    x, y = datasets.load_iris(return_X_y=True)
+    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.33, random_state=42)
+
+    # Lookup with a string
+    classifier = classifier_resolver.make("LogisticRegression")
+    classifier.fit(x_train, y_train)
+    assert 0.7 < classifier.score(x_test, y_test)
+
+    # Default lookup gives logistic regression
+    classifier = classifier_resolver.make(None)
+    classifier.fit(x_train, y_train)
+    assert 0.7 < classifier.score(x_test, y_test)
+
+.. seealso:: https://scikit-learn.org/stable/modules/classes.html#linear-classifiers
+"""

tests/test_contrib/test_sklearn.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+
+"""Tests for the scikit-learn contribution module."""
+
+import unittest
+
+try:
+    import sklearn
+except ImportError:  # pragma: no cover
+    sklearn = None  # pragma: no cover
+
+
+@unittest.skipUnless(sklearn, "Can not test sklearn contrib without ``pip install scikit-learn``.")
+class TestSklearn(unittest.TestCase):
+    """Test for the scikit-learn contribution module."""
+
+    def test_classifier_resolver(self):
+        """Tests for the classifier resolver."""
+        from sklearn import datasets
+        from sklearn.model_selection import train_test_split
+
+        from class_resolver.contrib.sklearn import classifier_resolver
+
+        x, y = datasets.load_iris(return_X_y=True)
+        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.33, random_state=42)
+
+        classifier = classifier_resolver.make("LogisticRegression")
+        classifier.fit(x_train, y_train)
+        accuracy = classifier.score(x_test, y_test)
+        self.assertLessEqual(0.0, accuracy)
+        self.assertGreaterEqual(1.0, accuracy)
+
+        for name, cls in classifier_resolver.lookup_dict.items():
+            with self.subTest(name=name):
+                classifier = cls()
+                classifier.fit(x_train, y_train)
+                accuracy = classifier.score(x_test, y_test)
+                self.assertLessEqual(0.0, accuracy)
+                self.assertGreaterEqual(1.0, accuracy)

tox.ini

@@ -36,6 +36,7 @@ extras =
     numpy
     optuna
     docdata
+    sklearn
 
 [testenv:coverage-clean]
 deps = coverage
@@ -114,6 +115,7 @@ extras =
     torch
     numpy
     optuna
+    sklearn
 commands =
     python -m sphinx -W -b html -d docs/build/doctrees docs/source docs/build/html