@Override
protected void moveOffsetAfter(boolean success) {
if (success && (myPanel != null && myPanel.getInitPlace() != InitPlace.SAME_METHOD)
|| myOperation.getInplaceInitPlace() != InitPlace.SAME_METHOD) {
final AccessToken accessToken =
ApplicationManager.getApplication().acquireWriteActionLock(getClass());
try {
final PyAssignmentStatement initializer =
PsiTreeUtil.getParentOfType(myTarget, PyAssignmentStatement.class);
assert initializer != null;
final Function<String, PyStatement> callback =
FunctionUtil.<String, PyStatement>constant(initializer);
final PyClass pyClass = PyUtil.getContainingClassOrSelf(initializer);
InitPlace initPlace =
myPanel != null ? myPanel.getInitPlace() : myOperation.getInplaceInitPlace();
if (initPlace == InitPlace.CONSTRUCTOR) {
AddFieldQuickFix.addFieldToInit(myProject, pyClass, "", callback);
} else if (initPlace == InitPlace.SET_UP) {
addFieldToSetUp(pyClass, callback);
}
if (myOperation.getOccurrences().size() > 0) {
initializer.delete();
} else {
final PyExpression copy =
PyElementGenerator.getInstance(myProject)
.createExpressionFromText(
LanguageLevel.forElement(myTarget), myTarget.getText());
initializer.replace(copy);
}
initializer.delete();
} finally {
accessToken.finish();
}
}
}
@Nullable
private static PsiElement addFieldToSetUp(
PyClass clazz, final Function<String, PyStatement> callback) {
final PyFunction init =
clazz.findMethodByName(PythonUnitTestUtil.TESTCASE_SETUP_NAME, false, null);
if (init != null) {
return AddFieldQuickFix.appendToMethod(init, callback);
}
final PyFunctionBuilder builder =
new PyFunctionBuilder(PythonUnitTestUtil.TESTCASE_SETUP_NAME, clazz);
builder.parameter(PyNames.CANONICAL_SELF);
PyFunction setUp = builder.buildFunction(clazz.getProject(), LanguageLevel.getDefault());
final PyStatementList statements = clazz.getStatementList();
final PsiElement anchor = statements.getFirstChild();
setUp = (PyFunction) statements.addBefore(setUp, anchor);
return AddFieldQuickFix.appendToMethod(setUp, callback);
}
@Nullable
@Override
protected PsiElement addDeclaration(
@NotNull PsiElement expression,
@NotNull PsiElement declaration,
@NotNull IntroduceOperation operation) {
final PsiElement expr = expression instanceof PyClass ? expression : expression.getParent();
PsiElement anchor = PyUtil.getContainingClassOrSelf(expr);
assert anchor instanceof PyClass;
final PyClass clazz = (PyClass) anchor;
final Project project = anchor.getProject();
if (operation.getInitPlace() == InitPlace.CONSTRUCTOR && !inConstructor(expression)) {
return AddFieldQuickFix.addFieldToInit(
project, clazz, "", new AddFieldDeclaration(declaration));
} else if (operation.getInitPlace() == InitPlace.SET_UP) {
return addFieldToSetUp(clazz, new AddFieldDeclaration(declaration));
}
return PyIntroduceVariableHandler.doIntroduceVariable(
expression, declaration, operation.getOccurrences(), operation.isReplaceAll());
}