Fill generics, when useful

apache · Nov 23, 2020 · 9182a98 · 9182a98
1 parent 426c7bb
commit 9182a98
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/python/pyspark/ml/feature.pyi b/python/pyspark/ml/feature.pyi
@@ -1602,7 +1602,10 @@ class _VarianceThresholdSelectorParams(HasFeaturesCol, HasOutputCol):
     def getVarianceThreshold(self) -> float: ...
 
 class VarianceThresholdSelector(
-    JavaEstimator, _VarianceThresholdSelectorParams, JavaMLReadable, JavaMLWritable
+    JavaEstimator[VarianceThresholdSelectorModel],
+    _VarianceThresholdSelectorParams,
+    JavaMLReadable[VarianceThresholdSelector],
+    JavaMLWritable,
 ):
     def __init__(
         self,
@@ -1621,7 +1624,10 @@ class VarianceThresholdSelector(
     def setOutputCol(self, value: str) -> VarianceThresholdSelector: ...
 
 class VarianceThresholdSelectorModel(
-    JavaModel, _VarianceThresholdSelectorParams, JavaMLReadable, JavaMLWritable
+    JavaModel,
+    _VarianceThresholdSelectorParams,
+    JavaMLReadable[VarianceThresholdSelectorModel],
+    JavaMLWritable,
 ):
     def setFeaturesCol(self, value: str) -> VarianceThresholdSelectorModel: ...
     def setOutputCol(self, value: str) -> VarianceThresholdSelectorModel: ...

diff --git a/python/pyspark/ml/pipeline.pyi b/python/pyspark/ml/pipeline.pyi
@@ -51,7 +51,7 @@ class PipelineWriter(MLWriter):
     def __init__(self, instance: Pipeline) -> None: ...
     def saveImpl(self, path: str) -> None: ...
 
-class PipelineReader(MLReader):
+class PipelineReader(MLReader[Pipeline]):
     cls: Type[Pipeline]
     def __init__(self, cls: Type[Pipeline]) -> None: ...
     def load(self, path: str) -> Pipeline: ...
@@ -61,7 +61,7 @@ class PipelineModelWriter(MLWriter):
     def __init__(self, instance: PipelineModel) -> None: ...
     def saveImpl(self, path: str) -> None: ...
 
-class PipelineModelReader(MLReader):
+class PipelineModelReader(MLReader[PipelineModel]):
     cls: Type[PipelineModel]
     def __init__(self, cls: Type[PipelineModel]) -> None: ...
     def load(self, path: str) -> PipelineModel: ...

diff --git a/python/pyspark/ml/regression.pyi b/python/pyspark/ml/regression.pyi
@@ -414,7 +414,7 @@ class RandomForestRegressionModel(
     _TreeEnsembleModel,
     _RandomForestRegressorParams,
     JavaMLWritable,
-    JavaMLReadable,
+    JavaMLReadable[RandomForestRegressionModel],
 ):
     @property
     def trees(self) -> List[DecisionTreeRegressionModel]: ...