com.jetbrains.python.psi.impl.blockEvaluator.PyBlockEvaluator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of python-community Show documentation
A packaging of the IntelliJ Community Edition python-community library. This is release number 1 of trunk branch 142.
The newest version!
/*
 * Copyright 2000-2014 JetBrains s.r.o.
 *
 * Licensed 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.
 */
package com.jetbrains.python.psi.impl.blockEvaluator;

import com.google.common.collect.Sets;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.intellij.psi.util.QualifiedName;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyEvaluator;
import com.jetbrains.python.psi.impl.PyPathEvaluator;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * @author yole
 */
public class PyBlockEvaluator {
  @NotNull
  private final PyEvaluationResult myEvaluationResult = new PyEvaluationResult();
  @NotNull
  private final PyEvaluationContext myContext;
  private final Set myVisitedFiles;
  private final Set myDeclarationsToTrack = new HashSet();
  private String myCurrentFilePath;
  private Object myReturnValue;
  private boolean myEvaluateCollectionItems = true;

  /**
   * @param evaluationContext context, obtained via {@link #getContext()}. Pass it here to enable cache. See {@link com.jetbrains.python.psi.impl.blockEvaluator.PyEvaluationContext}
   *                          for more info
   * @see com.jetbrains.python.psi.impl.blockEvaluator.PyEvaluationContext
   */
  public PyBlockEvaluator(@NotNull final PyEvaluationContext evaluationContext) {
    this(Sets.newHashSet(), evaluationContext);
  }

  /**
   * Create evaluator with out of cache context
   */
  public PyBlockEvaluator() {
    this(new PyEvaluationContext());
  }

  private PyBlockEvaluator(@NotNull final Set visitedFiles, @NotNull final PyEvaluationContext evaluationContext) {
    myVisitedFiles = visitedFiles;
    myContext = evaluationContext;
  }

  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 MyPyElementVisitor());
  }

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

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

    if (myDeclarationsToTrack.contains(nameBeingExtended)) {
      List declarations = myEvaluationResult.myDeclarations.get(nameBeingExtended);
      if (declarations != null) {
        PyPsiUtils.sequenceToList(declarations, arg);
      }
    }
  }

  private void processUpdateCall(PyCallExpression node, String name) {
    Object value = myEvaluationResult.myNamespace.get(name);
    if (value instanceof Map) {
      Object argValue = prepareEvaluator().evaluate(node.getArguments()[0]);
      if (argValue instanceof Map) {
        ((Map)value).putAll((Map)argValue);
      }
    }
  }

  private PyEvaluator prepareEvaluator() {
    PyEvaluator evaluator = createEvaluator();
    evaluator.setNamespace(myEvaluationResult.myNamespace);
    evaluator.setEvaluateCollectionItems(myEvaluateCollectionItems);
    return evaluator;
  }

  protected PyEvaluator createEvaluator() {
    return new PyPathEvaluator(myCurrentFilePath);
  }

  public Object getValue(String name) {
    return myEvaluationResult.myNamespace.get(name);
  }

  @Nullable
  public String getValueAsString(String name) {
    Object value = myEvaluationResult.myNamespace.get(name);
    return value instanceof String ? (String)value : null;
  }

  @Nullable
  public List getValueAsList(String name) {
    Object value = myEvaluationResult.myNamespace.get(name);
    return value instanceof List ? (List)value : null;
  }

  @NotNull
  public List getValueAsStringList(String name) {
    Object value = myEvaluationResult.myNamespace.get(name);
    if (value instanceof List) {
      List valueList = (List)value;
      for (Object o : valueList) {
        if (o != null && !(o instanceof String)) {
          return Collections.emptyList();
        }
      }
      return (List)value;
    }
    if (value instanceof String) {
      return Collections.singletonList((String)value);
    }
    return Collections.emptyList();
  }

  public Set getVisitedFiles() {
    return myVisitedFiles;
  }


  public Object getReturnValue() {
    return myReturnValue;
  }

  public void setEvaluateCollectionItems(boolean evaluateCollectionItems) {
    myEvaluateCollectionItems = evaluateCollectionItems;
  }

  @NotNull
  public List getDeclarations(@NotNull final String name) {
    return myEvaluationResult.getDeclarations(name);
  }

  /**
   * @return so-called context. You may pass it to any instance of {@link com.jetbrains.python.psi.impl.blockEvaluator.PyBlockEvaluator}
   * to make instances share their cache
   */
  @NotNull
  public PyEvaluationContext getContext() {
    return myContext;
  }

  private class MyPyElementVisitor extends PyElementVisitor {
    @Override
    public void visitPyAssignmentStatement(PyAssignmentStatement node) {
      PyExpression expression = node.getLeftHandSideExpression();
      if (expression instanceof PyTargetExpression) {
        String name = expression.getName();
        PyExpression value = ((PyTargetExpression)expression).findAssignedValue();
        myEvaluationResult.myNamespace.put(name, prepareEvaluator().evaluate(value));
        if (myDeclarationsToTrack.contains(name)) {
          List declarations = new ArrayList();
          PyPsiUtils.sequenceToList(declarations, value);
          myEvaluationResult.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 = myEvaluationResult.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).isQualified() && name != null) {
        Object currentValue = myEvaluationResult.myNamespace.get(name);
        if (currentValue != null) {
          Object rhs = prepareEvaluator().evaluate(node.getValue());
          myEvaluationResult.myNamespace.put(name, prepareEvaluator().concatenate(currentValue, rhs));
        }
        if (myDeclarationsToTrack.contains(name)) {
          List declarations = myEvaluationResult.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.isQualified()) {
            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(final PyFromImportStatement node) {
      if (node.isFromFuture()) return;
      final PsiElement source = PyUtil.turnDirIntoInit(node.resolveImportSource());
      if (source instanceof PyFile) {
        final PyFile pyFile = (PyFile)source;
        PyEvaluationResult newlyEvaluatedResult = myContext.getCachedResult(pyFile);

        if (newlyEvaluatedResult == null) {
          final PyBlockEvaluator importEvaluator = new PyBlockEvaluator(myVisitedFiles, myContext);
          importEvaluator.myDeclarationsToTrack.addAll(myDeclarationsToTrack);
          importEvaluator.evaluate(pyFile);
          newlyEvaluatedResult = importEvaluator.myEvaluationResult;
          myContext.cache(pyFile, newlyEvaluatedResult);
        }

        if (node.isStarImport()) {
          // TODO honor __all__ here
          myEvaluationResult.myNamespace.putAll(newlyEvaluatedResult.myNamespace);
          myEvaluationResult.myDeclarations.putAll(newlyEvaluatedResult.myDeclarations);
        }
        else {
          for (final PyImportElement element : node.getImportElements()) {
            final String nameOfVarInOurModule = element.getVisibleName();
            final QualifiedName nameOfVarInExternalModule = element.getImportedQName();
            if ((nameOfVarInOurModule == null) || (nameOfVarInExternalModule == null)) {
              continue;
            }

            final Object value = newlyEvaluatedResult.myNamespace.get(nameOfVarInExternalModule.toString());
            myEvaluationResult.myNamespace.put(nameOfVarInOurModule, value);
            final List declarations = newlyEvaluatedResult.getDeclarations(nameOfVarInOurModule);
            if (myEvaluationResult.myDeclarations.containsKey(nameOfVarInOurModule)) {
              myEvaluationResult.myDeclarations.get(nameOfVarInOurModule).addAll(declarations);
            }
            else {
              myEvaluationResult.myDeclarations.put(nameOfVarInOurModule, 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());
    }
  }
}