<commit_before>package com.jetbrains.python.psi.impl;

import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.psi.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * @author yole
 */
public class PyBlockEvaluator {
  private final Map<String, Object> myNamespace = new HashMap<String, Object>();
  private final Map<String, List<PyExpression>> myDeclarations = new HashMap<String, List<PyExpression>>();
  private final Set<PyFile> myVisitedFiles;
  private final Set<String> myDeclarationsToTrack = new HashSet<String>();
  private String myCurrentFilePath;
  private Object myReturnValue;
  private boolean myEvaluateCollectionItems = true;

  public PyBlockEvaluator() {
    myVisitedFiles = new HashSet<PyFile>();
  }

  public PyBlockEvaluator(Set<PyFile> visitedFiles) {
    myVisitedFiles = visitedFiles;
  }

  public void trackDeclarations(String attrName) {
    myDeclarationsToTrack.add(attrName);
  }

  public void evaluate(PyElement element) {
    VirtualFile vFile = element.getContainingFile().getVirtualFile();
    myCurrentFilePath = vFile != null ? vFile.getPath() : null;
    if (myVisitedFiles.contains(element)) {
      return;
    }
    myVisitedFiles.add((PyFile)element.getContainingFile());
    PyElement statementContainer = element instanceof PyFunction ? ((PyFunction) element).getStatementList() : element;
    if (statementContainer == null) {
      return;
    }
    statementContainer.acceptChildren(new PyElementVisitor() {
      @Override
      public void visitPyAssignmentStatement(PyAssignmentStatement node) {
        PyExpression expression = node.getLeftHandSideExpression();
        if (expression instanceof PyTargetExpression) {
          String name = expression.getName();
          PyExpression value = ((PyTargetExpression)expression).findAssignedValue();
          myNamespace.put(name, prepareEvaluator().evaluate(value));
          if (myDeclarationsToTrack.contains(name)) {
            List<PyExpression> declarations = new ArrayList<PyExpression>();
            PyPsiUtils.sequenceToList(declarations, value);
            myDeclarations.put(name, declarations);
          }
        }
        else if (expression instanceof PySubscriptionExpression) {
          PyExpression operand = ((PySubscriptionExpression)expression).getOperand();
          PyExpression indexExpression = ((PySubscriptionExpression)expression).getIndexExpression();
          if (operand instanceof PyReferenceExpression && ((PyReferenceExpression)operand).getQualifier() == null) {
            Object currentValue = myNamespace.get(((PyReferenceExpression)operand).getReferencedName());
            if (currentValue instanceof Map) {
              Object mapKey = prepareEvaluator().evaluate(indexExpression);
              if (mapKey != null) {
                Object value = myEvaluateCollectionItems ? prepareEvaluator().evaluate(node.getAssignedValue()) : node.getAssignedValue();
                ((Map)currentValue).put(mapKey, value);
              }
            }
          }
        }
      }

      @Override
      public void visitPyAugAssignmentStatement(PyAugAssignmentStatement node) {
        PyExpression target = node.getTarget();
        String name = target.getName();
        if (target instanceof PyReferenceExpression && ((PyReferenceExpression)target).getQualifier() == null && name != null) {
          Object currentValue = myNamespace.get(name);
          if (currentValue != null) {
            Object rhs = prepareEvaluator().evaluate(node.getValue());
            myNamespace.put(name, PyEvaluator.concatenate(currentValue, rhs));
          }
          if (myDeclarationsToTrack.contains(name)) {
            List<PyExpression> declarations = myDeclarations.get(name);
            if (declarations != null) {
              PyPsiUtils.sequenceToList(declarations, node.getValue());
            }
          }
        }
      }

      @Override
      public void visitPyExpressionStatement(PyExpressionStatement node) {
        node.getExpression().accept(this);
      }

      @Override
      public void visitPyCallExpression(PyCallExpression node) {
        PyExpression callee = node.getCallee();
        if (callee instanceof PyReferenceExpression) {
          PyReferenceExpression calleeRef = (PyReferenceExpression)callee;
          PyExpression qualifier = calleeRef.getQualifier();
          if (qualifier instanceof PyReferenceExpression) {
            PyReferenceExpression qualifierRef = (PyReferenceExpression)qualifier;
            if (qualifierRef.getQualifier() == null) {
              if (PyNames.EXTEND.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
                processExtendCall(node, qualifierRef.getReferencedName());
              }
              else if (PyNames.UPDATE.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
                processUpdateCall(node, qualifierRef.getReferencedName());
              }
            }
          }
        }
      }

      @Override
      public void visitPyFromImportStatement(PyFromImportStatement node) {
        if (node.isFromFuture()) return;
        PsiElement source = PyUtil.turnDirIntoInit(node.resolveImportSource());
        if (source instanceof PyFile) {
          PyBlockEvaluator importEvaluator = new PyBlockEvaluator(myVisitedFiles);
          importEvaluator.myDeclarationsToTrack.addAll(myDeclarationsToTrack);
          importEvaluator.evaluate((PyFile)source);
          if (node.isStarImport()) {
            // TODO honor __all__ here
            myNamespace.putAll(importEvaluator.myNamespace);
            myDeclarations.putAll(importEvaluator.myDeclarations);
          }
          else {
            for (PyImportElement element : node.getImportElements()) {
              Object value = importEvaluator.myNamespace.get(element.getName());
              String name = element.getAsName();
              if (name == null) {
                name = element.getName();
              }
              myNamespace.put(name, value);
              List<PyExpression> declarations = importEvaluator.getDeclarations(name);
              if (myDeclarations.containsKey(name)) {
                myDeclarations.get(name).addAll(declarations);
              }
              else {
                myDeclarations.put(name, declarations);
              }
            }
          }
        }
      }

      @Override
      public void visitPyReturnStatement(PyReturnStatement node) {
        myReturnValue = prepareEvaluator().evaluate(node.getExpression());
      }
    });
  }

  private void processExtendCall(PyCallExpression node, String nameBeingExtended) {
    PyExpression arg = node.getArguments()[0];

    Object value = myNamespace.get(nameBeingExtended);
    if (value instanceof List) {
      Object argValue = prepareEvaluator().evaluate(arg);
      myNamespace.put(nameBeingExtended, PyEvaluator.concatenate(value, argValue));
    }

    if (myDeclarationsToTrack.contains(nameBeingExtended)) {
      List<PyExpression> declarations = myDeclarations.get(nameBeingExtended);
      if (declarations != null) {
        PyPsiUtils.sequenceToList(declarations, arg);
      }
    }
  }<commit_msg>when parsing settings.py, go into true branch of if statements (PY-8493)<commit_after>package com.jetbrains.python.psi.impl;

import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.psi.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * @author yole
 */
public class PyBlockEvaluator {
  private final Map<String, Object> myNamespace = new HashMap<String, Object>();
  private final Map<String, List<PyExpression>> myDeclarations = new HashMap<String, List<PyExpression>>();
  private final Set<PyFile> myVisitedFiles;
  private final Set<String> myDeclarationsToTrack = new HashSet<String>();
  private String myCurrentFilePath;
  private Object myReturnValue;
  private boolean myEvaluateCollectionItems = true;

  public PyBlockEvaluator() {
    myVisitedFiles = new HashSet<PyFile>();
  }

  public PyBlockEvaluator(Set<PyFile> visitedFiles) {
    myVisitedFiles = visitedFiles;
  }

  public void trackDeclarations(String attrName) {
    myDeclarationsToTrack.add(attrName);
  }

  public void evaluate(PyElement element) {
    VirtualFile vFile = element.getContainingFile().getVirtualFile();
    myCurrentFilePath = vFile != null ? vFile.getPath() : null;
    if (myVisitedFiles.contains(element)) {
      return;
    }
    myVisitedFiles.add((PyFile)element.getContainingFile());
    PyElement statementContainer = element instanceof PyFunction ? ((PyFunction) element).getStatementList() : element;
    if (statementContainer == null) {
      return;
    }
    statementContainer.acceptChildren(new PyElementVisitor() {
      @Override
      public void visitPyAssignmentStatement(PyAssignmentStatement node) {
        PyExpression expression = node.getLeftHandSideExpression();
        if (expression instanceof PyTargetExpression) {
          String name = expression.getName();
          PyExpression value = ((PyTargetExpression)expression).findAssignedValue();
          myNamespace.put(name, prepareEvaluator().evaluate(value));
          if (myDeclarationsToTrack.contains(name)) {
            List<PyExpression> declarations = new ArrayList<PyExpression>();
            PyPsiUtils.sequenceToList(declarations, value);
            myDeclarations.put(name, declarations);
          }
        }
        else if (expression instanceof PySubscriptionExpression) {
          PyExpression operand = ((PySubscriptionExpression)expression).getOperand();
          PyExpression indexExpression = ((PySubscriptionExpression)expression).getIndexExpression();
          if (operand instanceof PyReferenceExpression && ((PyReferenceExpression)operand).getQualifier() == null) {
            Object currentValue = myNamespace.get(((PyReferenceExpression)operand).getReferencedName());
            if (currentValue instanceof Map) {
              Object mapKey = prepareEvaluator().evaluate(indexExpression);
              if (mapKey != null) {
                Object value = myEvaluateCollectionItems ? prepareEvaluator().evaluate(node.getAssignedValue()) : node.getAssignedValue();
                ((Map)currentValue).put(mapKey, value);
              }
            }
          }
        }
      }

      @Override
      public void visitPyAugAssignmentStatement(PyAugAssignmentStatement node) {
        PyExpression target = node.getTarget();
        String name = target.getName();
        if (target instanceof PyReferenceExpression && ((PyReferenceExpression)target).getQualifier() == null && name != null) {
          Object currentValue = myNamespace.get(name);
          if (currentValue != null) {
            Object rhs = prepareEvaluator().evaluate(node.getValue());
            myNamespace.put(name, PyEvaluator.concatenate(currentValue, rhs));
          }
          if (myDeclarationsToTrack.contains(name)) {
            List<PyExpression> declarations = myDeclarations.get(name);
            if (declarations != null) {
              PyPsiUtils.sequenceToList(declarations, node.getValue());
            }
          }
        }
      }

      @Override
      public void visitPyExpressionStatement(PyExpressionStatement node) {
        node.getExpression().accept(this);
      }

      @Override
      public void visitPyCallExpression(PyCallExpression node) {
        PyExpression callee = node.getCallee();
        if (callee instanceof PyReferenceExpression) {
          PyReferenceExpression calleeRef = (PyReferenceExpression)callee;
          PyExpression qualifier = calleeRef.getQualifier();
          if (qualifier instanceof PyReferenceExpression) {
            PyReferenceExpression qualifierRef = (PyReferenceExpression)qualifier;
            if (qualifierRef.getQualifier() == null) {
              if (PyNames.EXTEND.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
                processExtendCall(node, qualifierRef.getReferencedName());
              }
              else if (PyNames.UPDATE.equals(calleeRef.getReferencedName()) && node.getArguments().length == 1) {
                processUpdateCall(node, qualifierRef.getReferencedName());
              }
            }
          }
        }
      }

      @Override
      public void visitPyFromImportStatement(PyFromImportStatement node) {
        if (node.isFromFuture()) return;
        PsiElement source = PyUtil.turnDirIntoInit(node.resolveImportSource());
        if (source instanceof PyFile) {
          PyBlockEvaluator importEvaluator = new PyBlockEvaluator(myVisitedFiles);
          importEvaluator.myDeclarationsToTrack.addAll(myDeclarationsToTrack);
          importEvaluator.evaluate((PyFile)source);
          if (node.isStarImport()) {
            // TODO honor __all__ here
            myNamespace.putAll(importEvaluator.myNamespace);
            myDeclarations.putAll(importEvaluator.myDeclarations);
          }
          else {
            for (PyImportElement element : node.getImportElements()) {
              Object value = importEvaluator.myNamespace.get(element.getName());
              String name = element.getAsName();
              if (name == null) {
                name = element.getName();
              }
              myNamespace.put(name, value);
              List<PyExpression> declarations = importEvaluator.getDeclarations(name);
              if (myDeclarations.containsKey(name)) {
                myDeclarations.get(name).addAll(declarations);
              }
              else {
                myDeclarations.put(name, declarations);
              }
            }
          }
        }
      }

      @Override
      public void visitPyIfStatement(PyIfStatement node) {
        PyStatementList list = node.getIfPart().getStatementList();
        if (list != null) {
          list.acceptChildren(this);
        }
      }

      @Override
      public void visitPyReturnStatement(PyReturnStatement node) {
        myReturnValue = prepareEvaluator().evaluate(node.getExpression());
      }
    });
  }

  private void processExtendCall(PyCallExpression node, String nameBeingExtended) {
    PyExpression arg = node.getArguments()[0];

    Object value = myNamespace.get(nameBeingExtended);
    if (value instanceof List) {
      Object argValue = prepareEvaluator().evaluate(arg);
      myNamespace.put(nameBeingExtended, PyEvaluator.concatenate(value, argValue));
    }

    if (myDeclarationsToTrack.contains(nameBeingExtended)) {
      List<PyExpression> declarations = myDeclarations.get(nameBeingExtended);
      if (declarations != null) {
        PyPsiUtils.sequenceToList(declarations, arg);
      }
    }
  }
{'fraction_non_alphanumeric': 0.09891870244293152, 'fraction_numerical': 0.0006674676278200508, 'mean_word_length': 2.8806008806008805, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 35, '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': '6719339', 'n_tokens_mistral': 3937, 'n_tokens_neox': 3836, 'n_words': 726}