/** * Creates abstract version of each method in each class (does not touch method itself as opposite * to {@link #moveMethods(com.jetbrains.python.psi.PyClass, java.util.Collection, * com.jetbrains.python.psi.PyClass...)}) * * @param currentFunctions functions to make them abstract * @param to classes where abstract method should be created */ private static void makeMethodsAbstract( final Collection<PyFunction> currentFunctions, final PyClass... to) { final Set<PsiFile> filesToCheckImport = new HashSet<PsiFile>(); final Set<PyClass> classesToAddMetaAbc = new HashSet<PyClass>(); for (final PyFunction function : currentFunctions) { for (final PyClass destClass : to) { final PyFunctionBuilder functionBuilder = PyFunctionBuilder.copySignature(function, DECORATORS_MAY_BE_COPIED_TO_ABSTRACT); functionBuilder.decorate(PyNames.ABSTRACTMETHOD); final LanguageLevel level = LanguageLevel.forElement(destClass); PyClassRefactoringUtil.addMethods( destClass, false, functionBuilder.buildFunction(destClass.getProject(), level)); classesToAddMetaAbc.add(destClass); } } // Add ABCMeta to new classes if needed for (final PyClass aClass : classesToAddMetaAbc) { if (addMetaAbcIfNeeded(aClass)) { filesToCheckImport.add(aClass.getContainingFile()); } } // Add imports for ABC if needed for (final PsiFile file : filesToCheckImport) { addImportFromAbc(file, PyNames.ABSTRACTMETHOD); addImportFromAbc(file, ABC_META_CLASS); PyClassRefactoringUtil.optimizeImports(file); // To remove redundant imports } }
@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); }