plugins/beam-code-completion-plugin/src/main/java/BeamCompletionContributor.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import com.intellij.codeInsight.completion.*;

import com.intellij.patterns.PatternCondition;
import com.intellij.patterns.PsiJavaElementPattern;
import com.intellij.patterns.PsiJavaPatterns;
import static com.intellij.patterns.PlatformPatterns.psiElement;

import com.intellij.psi.*;

import com.intellij.codeInsight.lookup.LookupElementBuilder;
import com.intellij.util.ProcessingContext;
import org.jetbrains.annotations.NotNull;

/**
 * The `BeamCompletionContributor` class is a subclass of `CompletionContributor` that provides code completion
 * suggestions specifically for Apache Beam pipelines.
 */
public class BeamCompletionContributor extends CompletionContributor {
    public static final String[] beamJavaSDKTransforms = {
        "Filter", "FlatMapElements", "Keys", "KvSwap", "MapElements", "ParDo",
                "Partition", "Regex", "Reify", "ToString", "WithKeys", "WithTimestamps",
                "Values", "ApproximateQuantiles", "ApproximateUnique", "CoGroupByKey", "Combine",
                "CombineWithContext", "Count", "Distinct", "GroupByKey", "GroupIntoBatches",
                "HllCount", "Latest", "Max", "Mean", "Min", "Sample", "Sum", "Top",
                "Create", "Flatten", "PAssert", "View", "Window"
    };

    /**
     * A pattern condition that matches method call expressions with the name "apply" in the context of the
     * Apache Beam `Pipeline` class.
     */
    public static final PatternCondition<PsiMethodCallExpression> APPLY_METHOD_PATTERN = new PatternCondition<>("") {
        @Override
        public boolean accepts(@NotNull PsiMethodCallExpression psiMethodCallExpression, ProcessingContext context) {
            String referenceName = psiMethodCallExpression.getMethodExpression().getReferenceName();
            if (referenceName != null){
                if (!referenceName.equals("apply")) {
                    return false;
                }
                PsiMethod resolvedMethod = psiMethodCallExpression.resolveMethod();
                if (resolvedMethod != null){
                    PsiClass containingClass = resolvedMethod.getContainingClass();
                    if (containingClass != null){
                        return "org.apache.beam.sdk.Pipeline".equals(containingClass.getQualifiedName()) ||
                                "org.apache.beam.sdk.values.PCollection".equals(containingClass.getQualifiedName());
                    }
                }
            }
            return false;
        }
    };

    /**
     * A pattern that matches an identifier after a method call expression that satisfies the `APPLY_METHOD_PATTERN`.
     */
    private static final PsiJavaElementPattern.Capture<PsiIdentifier> AFTER_METHOD_CALL_PATTERN =
            PsiJavaPatterns.psiElement(PsiIdentifier.class)
                    .withParent(
                            psiElement(PsiReferenceExpression.class)
                                    .withParent(
                                            psiElement(PsiExpressionList.class)
                                                    .withParent(
                                                            psiElement(PsiMethodCallExpression.class)
                                                                    .with(APPLY_METHOD_PATTERN)
                                                    )
                                    )
                    );

    /**
     * Fills the completion variants with Transforms from Java SDK.
     */
    @Override
    public void fillCompletionVariants(@NotNull final CompletionParameters parameters, @NotNull CompletionResultSet result) {
        if (AFTER_METHOD_CALL_PATTERN.accepts(parameters.getPosition())){
            for (String transform: beamJavaSDKTransforms) {
                result.addElement(LookupElementBuilder.create(transform).appendTailText(" org.apache.beam.sdk.transforms", true));
            }
        }
    }
}