Merge pull request #4356 from vortex-ape/dict_vectorizer

amueller · amueller · commit 6e540793930d · 2015-04-01T16:44:40.000-05:00
Fixes #4355: DictVectorizer.restrict docstring unclear
diff --git a/sklearn/feature_extraction/dict_vectorizer.py b/sklearn/feature_extraction/dict_vectorizer.py
@@ -312,15 +312,37 @@ def get_feature_names(self):
         return self.feature_names_
 
     def restrict(self, support, indices=False):
-        """Restrict the features to those in support.
+        """Restrict the features to those in support using feature selection.
+
+        This function modifies the estimator in-place.
 
         Parameters
         ----------
         support : array-like
             Boolean mask or list of indices (as returned by the get_support
             member of feature selectors).
         indices : boolean, optional
-            Whether support is a list of indices.
+            Whether support is a list of indices. 
+
+        Returns
+        -------
+        self
+
+        Examples
+        --------
+        >>> from sklearn.feature_extraction import DictVectorizer
+        >>> from sklearn.feature_selection import SelectKBest, chi2
+        >>> v = DictVectorizer()
+        >>> D = [{'foo': 1, 'bar': 2}, {'foo': 3, 'baz': 1}]
+        >>> X = v.fit_transform(D)
+        >>> support = SelectKBest(chi2, k=2).fit(X, [0, 1])
+        >>> v.get_feature_names()
+        ['bar', 'baz', 'foo']
+        >>> v.restrict(support.get_support()) # doctest: +ELLIPSIS
+        DictVectorizer(dtype=..., separator='=', sort=True,
+                sparse=True)
+        >>> v.get_feature_names()
+        ['bar', 'foo']
         """
         if not indices:
             support = np.where(support)[0]
