<commit_before>package com.jetbrains.python.codeInsight;

import com.intellij.codeInsight.completion.*;
import com.intellij.codeInsight.lookup.LookupElement;
import com.intellij.codeInsight.lookup.LookupElementBuilder;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiReference;
import com.intellij.psi.filters.position.FilterPattern;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.Icons;
import com.intellij.util.ProcessingContext;
import com.jetbrains.python.PythonLanguage;
import com.jetbrains.python.psi.*;
import org.jetbrains.annotations.NotNull;


import static com.intellij.patterns.PlatformPatterns.psiElement;

/**
 * User: catherine
 *
 * Complete known keys for dictionaries
 */
public class PyDictKeyNamesCompletionContributor extends PySeeingOriginalCompletionContributor {
  @Override
  public AutoCompletionDecision handleAutoCompletionPossibility(AutoCompletionContext context) {
    // auto-insert the obvious only case; else show other cases.
    final LookupElement[] items = context.getItems();
    if (items.length == 1) {
      return AutoCompletionDecision.insertItem(items[0]);
    }
    return AutoCompletionDecision.SHOW_LOOKUP;
  }

  private static final FilterPattern DICT_KEY = new FilterPattern(
    new InSequenceFilter(psiElement(PySubscriptionExpression.class))
  );

  public PyDictKeyNamesCompletionContributor() {
    extend(
      CompletionType.BASIC,
      psiElement()
        .withLanguage(PythonLanguage.getInstance())
        .and(DICT_KEY)
     ,
      new CompletionProvider<CompletionParameters>() {
        protected void addCompletions(
          @NotNull final CompletionParameters parameters, final ProcessingContext context, @NotNull final CompletionResultSet result
        ) {
          PsiElement original = parameters.getOriginalPosition();
          PsiElement operand = PsiTreeUtil.getParentOfType(original, PySubscriptionExpression.class).getOperand();
          if (operand != null) {
            PsiReference reference = operand.getReference();
            if (reference != null) {
              PsiElement resolvedElement = reference.resolve();
              if (resolvedElement instanceof PyTargetExpression) {
                PyDictLiteralExpression dict = PsiTreeUtil.getNextSiblingOfType(resolvedElement, PyDictLiteralExpression.class);
                if (dict != null) {
                  addDictLiteralKeys(dict, result);
                  PsiFile file = parameters.getOriginalFile();
                  addAdditionalKeys(file, operand, result);
                }
              }
            }
          }
        }

        private void addAdditionalKeys(PsiFile file, PsiElement operand, CompletionResultSet result) {
          PySubscriptionExpression[] subscriptionExpressions = PyUtil.getAllChildrenOfType(file, PySubscriptionExpression.class);
          for (PySubscriptionExpression expr : subscriptionExpressions) {
            if (expr.getOperand().getText().equals(operand.getText())) {
              PsiElement parent = expr.getParent();
              if (parent instanceof PyAssignmentStatement) {
                if (expr.equals(((PyAssignmentStatement)parent).getLeftHandSideExpression())) {
                  PyExpression key = expr.getIndexExpression();
                  if (key != null) {
                    LookupElementBuilder item;
                    item = LookupElementBuilder
                      .create(key.getText())
                      .setTypeText("dict key")
                      .setIcon(Icons.PARAMETER_ICON);
                    result.addElement(item);
                  }
                }
              }
            }
          }
        }

        private void addDictLiteralKeys(PyDictLiteralExpression dict, CompletionResultSet result) {
          PyKeyValueExpression[] keyValues = dict.getElements();
          for (PyKeyValueExpression expression : keyValues) {
            LookupElementBuilder item;
            item = LookupElementBuilder
              .create(expression.getKey().getText())
              .setTypeText("dict key")
              .setIcon(Icons.PARAMETER_ICON);
            result.addElement(item);
          }
        }
      }
    );
  }
}
<commit_msg>added dict keys from dictionary constructor to dictNamesCompletion<commit_after>package com.jetbrains.python.codeInsight;

import com.intellij.codeInsight.completion.*;
import com.intellij.codeInsight.lookup.LookupElement;
import com.intellij.codeInsight.lookup.LookupElementBuilder;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiReference;
import com.intellij.psi.filters.position.FilterPattern;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.Icons;
import com.intellij.util.ProcessingContext;
import com.jetbrains.python.PythonLanguage;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.types.PyType;
import com.jetbrains.python.psi.types.TypeEvalContext;
import org.jetbrains.annotations.NotNull;


import static com.intellij.patterns.PlatformPatterns.psiElement;

/**
 * User: catherine
 *
 * Complete known keys for dictionaries
 */
public class PyDictKeyNamesCompletionContributor extends PySeeingOriginalCompletionContributor {
  @Override
  public AutoCompletionDecision handleAutoCompletionPossibility(AutoCompletionContext context) {
    // auto-insert the obvious only case; else show other cases.
    final LookupElement[] items = context.getItems();
    if (items.length == 1) {
      return AutoCompletionDecision.insertItem(items[0]);
    }
    return AutoCompletionDecision.SHOW_LOOKUP;
  }

  private static final FilterPattern DICT_KEY = new FilterPattern(
    new InSequenceFilter(psiElement(PySubscriptionExpression.class))
  );

  public PyDictKeyNamesCompletionContributor() {
    extend(
      CompletionType.BASIC,
      psiElement()
        .withLanguage(PythonLanguage.getInstance())
        .and(DICT_KEY)
     ,
      new CompletionProvider<CompletionParameters>() {
        protected void addCompletions(
          @NotNull final CompletionParameters parameters, final ProcessingContext context, @NotNull final CompletionResultSet result
        ) {
          PsiElement original = parameters.getOriginalPosition();
          PsiElement operand = PsiTreeUtil.getParentOfType(original, PySubscriptionExpression.class).getOperand();
          if (operand != null) {
            PsiReference reference = operand.getReference();
            if (reference != null) {
              PsiElement resolvedElement = reference.resolve();
              if (resolvedElement instanceof PyTargetExpression) {
                PyDictLiteralExpression dict = PsiTreeUtil.getNextSiblingOfType(resolvedElement, PyDictLiteralExpression.class);
                if (dict != null) {
                  addDictLiteralKeys(dict, result);
                  PsiFile file = parameters.getOriginalFile();
                  addAdditionalKeys(file, operand, result);
                }
                PyCallExpression dictConstructor = PsiTreeUtil.getNextSiblingOfType(resolvedElement, PyCallExpression.class);
                if (dictConstructor != null) {
                  addDictConstructorKeys(dictConstructor, result);
                  PsiFile file = parameters.getOriginalFile();
                  addAdditionalKeys(file, operand, result);
                }
              }
            }
          }
        }

        /**
         * add keys to completion result from dict constructor
         */
        private void addDictConstructorKeys(PyCallExpression dictConstructor, CompletionResultSet result) {
          String name = dictConstructor.getCallee().getText();
          if ("dict".equals(name)) {
            PyType type = dictConstructor.getType(TypeEvalContext.fast());
            if (type != null) {
              if (type.isBuiltin()) {
                PyExpression[] argumentList = dictConstructor.getArgumentList().getArguments();
                for (PyExpression argument : argumentList) {
                  if (argument instanceof PyKeywordArgument) {
                    LookupElementBuilder item;
                    item = LookupElementBuilder
                      .create("'" + ((PyKeywordArgument)argument).getKeyword() + "'")
                      .setTypeText("dict key")
                      .setIcon(Icons.PARAMETER_ICON);
                    result.addElement(item);
                  }
                }
              }
            }
          }
        }

        /**
         * add keys from assignment statements
         * For instance, dictionary['b']=b
         * @param file to get additional keys
         * @param operand is operand of origin element
         * @param result is completion result set
         */
        private void addAdditionalKeys(PsiFile file, PsiElement operand, CompletionResultSet result) {
          PySubscriptionExpression[] subscriptionExpressions = PyUtil.getAllChildrenOfType(file, PySubscriptionExpression.class);
          for (PySubscriptionExpression expr : subscriptionExpressions) {
            if (expr.getOperand().getText().equals(operand.getText())) {
              PsiElement parent = expr.getParent();
              if (parent instanceof PyAssignmentStatement) {
                if (expr.equals(((PyAssignmentStatement)parent).getLeftHandSideExpression())) {
                  PyExpression key = expr.getIndexExpression();
                  if (key != null) {
                    LookupElementBuilder item;
                    item = LookupElementBuilder
                      .create(key.getText())
                      .setTypeText("dict key")
                      .setIcon(Icons.PARAMETER_ICON);
                    result.addElement(item);
                  }
                }
              }
            }
          }
        }

        /**
         * add keys from dict literal expression
         */
        private void addDictLiteralKeys(PyDictLiteralExpression dict, CompletionResultSet result) {
          PyKeyValueExpression[] keyValues = dict.getElements();
          for (PyKeyValueExpression expression : keyValues) {
            LookupElementBuilder item;
            item = LookupElementBuilder
              .create(expression.getKey().getText())
              .setTypeText("dict key")
              .setIcon(Icons.PARAMETER_ICON);
            result.addElement(item);
          }
        }
      }
    );
  }
}

{'fraction_non_alphanumeric': 0.08504538939321549, 'fraction_numerical': 0.0003822264691829909, 'mean_word_length': 2.3991555699902567, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7395605', 'n_tokens_mistral': 2596, 'n_tokens_neox': 2526, 'n_words': 511}