private static boolean inConstructor(@NotNull PsiElement expression) {
final PsiElement expr = expression instanceof PyClass ? expression : expression.getParent();
PyClass clazz = PyUtil.getContainingClassOrSelf(expr);
final ScopeOwner current = ScopeUtil.getScopeOwner(expression);
if (clazz != null && current != null && current instanceof PyFunction) {
PyFunction init = clazz.findMethodByName(PyNames.INIT, false, null);
if (current == init) {
return true;
}
}
return false;
}
@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);
}
