SONARPY-1834 check if pipeline is used by sklearn.compose estimator (#2034)

* SONARPY-1834 check if pipeline is used by sklearn.compose estimator

* SONARPY-1834 small refactoring

Author: Seppli11

python-checks/src/main/java/org/sonar/python/checks/SklearnPipelineSpecifyMemoryArgumentCheck.java

@@ -56,7 +56,11 @@ private static void checkCallExpression(SubscriptionContext subscriptionContext)
             return;
           }
 
-          if (getAssignedName(callExpression).map(SklearnPipelineSpecifyMemoryArgumentCheck::isUsedInAnotherPipeline).orElse(false)) {
+          boolean isUsedInAnotherPipeline = getAssignedName(callExpression)
+            .map(SklearnPipelineSpecifyMemoryArgumentCheck::isUsedInAnotherPipeline)
+            .orElse(false);
+
+          if (isUsedInAnotherPipeline) {
             return;
           }
 
@@ -72,26 +76,44 @@ private static void createIssue(SubscriptionContext subscriptionContext, CallExp
     issue.addQuickFix(quickFix);
   }
 
-  private static boolean isPipelineCreation(CallExpression callExpression) {
-    return Optional.ofNullable(callExpression.calleeSymbol())
-      .map(Symbol::fullyQualifiedName)
-      .map(fqn -> "sklearn.pipeline.Pipeline".equals(fqn) || "sklearn.pipeline.make_pipeline".equals(fqn))
-      .orElse(false);
-  }
-
   private static boolean isUsedInAnotherPipeline(Name name) {
     Symbol symbol = name.symbol();
     return symbol != null && symbol.usages().stream().filter(usage -> !usage.isBindingUsage()).anyMatch(u -> {
       Tree tree = u.tree();
       CallExpression callExpression = (CallExpression) TreeUtils.firstAncestorOfKind(tree, Tree.Kind.CALL_EXPR);
       while (callExpression != null) {
-        Optional<String> fullyQualifiedName = Optional.ofNullable(callExpression.calleeSymbol()).map(Symbol::fullyQualifiedName);
-        if (fullyQualifiedName.isPresent() && isPipelineCreation(callExpression)) {
+        if (isUsedBySklearnComposeEstimatorOrPipelineCreation(callExpression)) {
           return true;
         }
         callExpression = (CallExpression) TreeUtils.firstAncestorOfKind(callExpression, Tree.Kind.CALL_EXPR);
       }
       return false;
     });
   }
+
+  private static boolean isPipelineCreation(CallExpression callExpression) {
+    return Optional.ofNullable(callExpression.calleeSymbol())
+      .map(Symbol::fullyQualifiedName)
+      .map(SklearnPipelineSpecifyMemoryArgumentCheck::isFullyQualifiedNameAPipelineCreation)
+      .orElse(false);
+  }
+
+  private static boolean isUsedBySklearnComposeEstimatorOrPipelineCreation(CallExpression callExpression) {
+    Symbol calleeSymbol = callExpression.calleeSymbol();
+    if(calleeSymbol == null) return false;
+
+    String fqn = calleeSymbol.fullyQualifiedName();
+    if(fqn == null) return false;
+
+    return isFullyQualifiedNameAPipelineCreation(fqn) || isFullyQualifiedNameASklearnComposeEstimator(fqn);
+  }
+
+  private static boolean isFullyQualifiedNameAPipelineCreation(String fqn) {
+    return "sklearn.pipeline.Pipeline".equals(fqn) || "sklearn.pipeline.make_pipeline".equals(fqn);
+  }
+
+  private static boolean isFullyQualifiedNameASklearnComposeEstimator(String fqn) {
+    return fqn.startsWith("sklearn.compose.");
+  }
+
 }

python-checks/src/test/resources/checks/sklearn_pipeline_specify_memory_argument.py

@@ -120,6 +120,17 @@ def __init__(self):
                                                        #^^^^^^^^^^^^^
             r1 = make_pipeline(p5) # Noncompliant
 
+def more_nested5():
+    from sklearn.pipeline import Pipeline
+    from sklearn.compose import ColumnTransformer, make_column_selector
+    numeric_transformer = Pipeline(steps=[('scaler', StandardScaler())])
+
+    preprocessor = ColumnTransformer(transformers=[('num', numeric_transformer, list(X_train.columns.values))])
+    clf = Pipeline(steps=[('preprocessor', preprocessor),
+                          ('classifier', LogisticRegression())])
+
+    preprocessor = make_column_selector(transformers=[('num', clf, list(X_train.columns.values))])
+
 
 def other():
     from sklearn.pipeline import Pipeline, make_pipeline