diff --git a/src/de/endrullis/idea/postfixtemplates/languages/python/CustomPythonStringPostfixTemplate.java b/src/de/endrullis/idea/postfixtemplates/languages/python/CustomPythonStringPostfixTemplate.java index 944bb7f6..2f7d56ac 100644 --- a/src/de/endrullis/idea/postfixtemplates/languages/python/CustomPythonStringPostfixTemplate.java +++ b/src/de/endrullis/idea/postfixtemplates/languages/python/CustomPythonStringPostfixTemplate.java @@ -1,59 +1,93 @@ package de.endrullis.idea.postfixtemplates.languages.python; +import java.util.HashMap; +import java.util.Map; + +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + import com.intellij.codeInsight.template.postfix.templates.PostfixTemplateProvider; import com.intellij.openapi.util.Condition; import com.intellij.psi.PsiElement; -import com.jetbrains.python.PyElementTypes; import com.jetbrains.python.codeInsight.postfix.PyPostfixUtils; import com.jetbrains.python.psi.PyTypedElement; +import com.jetbrains.python.psi.types.PyClassLikeType; import com.jetbrains.python.psi.types.PyType; -import com.jetbrains.python.psi.types.PyTypeTokenTypes; import com.jetbrains.python.psi.types.TypeEvalContext; + import de.endrullis.idea.postfixtemplates.templates.SimpleStringBasedPostfixTemplate; import de.endrullis.idea.postfixtemplates.templates.SpecialType; -import org.jetbrains.annotations.NotNull; -import org.jetbrains.annotations.Nullable; - -import java.util.HashMap; -import java.util.Map; /** * Custom postfix template for Python. */ @SuppressWarnings("WeakerAccess") public class CustomPythonStringPostfixTemplate extends SimpleStringBasedPostfixTemplate { + /** + * Contains predefined type-to-psiCondition mappings as well as cached mappings for individual types. + */ + private static final Map> type2psiCondition = + new HashMap>() {{ + put(SpecialType.ANY.name(), e -> true); - /** Contains predefined type-to-psiCondition mappings as well as cached mappings for individual types. */ - private static final Map> type2psiCondition = new HashMap>() {{ - put(SpecialType.ANY.name(), e -> true); - for (String pyType : PythonPostfixTemplatesUtils.PYTHON_TYPES) { - put(pyType, e -> { - if (e instanceof PyTypedElement) { - PyType type = TypeEvalContext.codeAnalysis(e.getProject(), e.getContainingFile()).getType((PyTypedElement) e); - return type != null && type.getName() != null && type.getName().equals(pyType); - } else { - return false; - } + PythonPostfixTemplatesUtils.COMPLEX_TYPES.forEach((key, value) -> { + put(key, e -> { + String typeName = getTypeName(e); + return value.contains(typeName); + }); }); - } - }}; + for (String pyType : PythonPostfixTemplatesUtils.PYTHON_TYPES) { + put(pyType, e -> { + if (e instanceof PyTypedElement) { + String typeName = getTypeName(e); + return pyType.equals(typeName); + } else { + return false; + } + }); + } + }}; + + @Nullable + private static PyType getType(PsiElement e) { + return TypeEvalContext.codeAnalysis(e.getProject(), e.getContainingFile()).getType((PyTypedElement) e); + } + + @Nullable + private static String getTypeName(PsiElement e) { + PyType type = getType(e); + final String name; + if (type != null) { + if (type instanceof PyClassLikeType) { + name = ((PyClassLikeType) type).getClassQName(); + } else { + name = type.getName(); + } + } else { + name = null; + } + return name; + } - public CustomPythonStringPostfixTemplate(String matchingClass, String conditionClass, String name, String example, String template, PostfixTemplateProvider provider, PsiElement psiElement) { - super(name, example, template, provider, psiElement, PyPostfixUtils.selectorAllExpressionsWithCurrentOffset(getCondition(matchingClass, conditionClass))); - } + public CustomPythonStringPostfixTemplate(String matchingClass, String conditionClass, String name, String example, + String template, PostfixTemplateProvider provider, PsiElement psiElement) { + super(name, example, template, provider, psiElement, + PyPostfixUtils.selectorAllExpressionsWithCurrentOffset(getCondition(matchingClass, conditionClass))); + } - @NotNull - public static Condition getCondition(final @NotNull String matchingClass, final @Nullable String conditionClass) { - Condition psiElementCondition = type2psiCondition.get(matchingClass); + @NotNull + public static Condition getCondition(final @NotNull String matchingClass, + final @Nullable String conditionClass) { + Condition psiElementCondition = type2psiCondition.get(matchingClass); - // PyElementTypes.INTEGER_LITERAL_EXPRESSION - //TypeEvalContext.codeAnalysis(e.getProject(), e.getContainingFile()).getType((PyTypedElement) e) + // PyElementTypes.INTEGER_LITERAL_EXPRESSION + //TypeEvalContext.codeAnalysis(e.getProject(), e.getContainingFile()).getType((PyTypedElement) e) - if (psiElementCondition == null) { - //psiElementCondition = PythonPostfixTemplatesUtils.isCustomClass(matchingClass); - } + if (psiElementCondition == null) { + //psiElementCondition = PythonPostfixTemplatesUtils.isCustomClass(matchingClass); + } - return psiElementCondition; - } + return psiElementCondition; + } } diff --git a/src/de/endrullis/idea/postfixtemplates/languages/python/PythonAnnotator.java b/src/de/endrullis/idea/postfixtemplates/languages/python/PythonAnnotator.java index 0f3de63b..5f5f369e 100644 --- a/src/de/endrullis/idea/postfixtemplates/languages/python/PythonAnnotator.java +++ b/src/de/endrullis/idea/postfixtemplates/languages/python/PythonAnnotator.java @@ -20,6 +20,9 @@ public class PythonAnnotator implements CptLangAnnotator { private final Map className2exists = new HashMap() {{ put(SpecialType.ANY.name(), true); + for (String complexType : PythonPostfixTemplatesUtils.COMPLEX_TYPES.keySet()) { + put(complexType, true); + } for (String pyType : PythonPostfixTemplatesUtils.PYTHON_TYPES) { put(pyType, true); } diff --git a/src/de/endrullis/idea/postfixtemplates/languages/python/PythonPostfixTemplatesUtils.java b/src/de/endrullis/idea/postfixtemplates/languages/python/PythonPostfixTemplatesUtils.java index 47a687c3..6f94b58e 100644 --- a/src/de/endrullis/idea/postfixtemplates/languages/python/PythonPostfixTemplatesUtils.java +++ b/src/de/endrullis/idea/postfixtemplates/languages/python/PythonPostfixTemplatesUtils.java @@ -1,9 +1,11 @@ package de.endrullis.idea.postfixtemplates.languages.python; -import java.util.Set; - import static de.endrullis.idea.postfixtemplates.utils.CollectionUtils._Set; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + /** * Utilities for Python postfix templates. * @@ -11,6 +13,22 @@ */ class PythonPostfixTemplatesUtils { + static final Map> COMPLEX_TYPES; + + static final Set ITERABLE_TYPES = _Set( + "list", + "dict", + "set", + "tuple", + "frozenset" + ); + + static { + COMPLEX_TYPES = new HashMap<>(); + // iter + COMPLEX_TYPES.put("iter", ITERABLE_TYPES); + } + static final Set PYTHON_TYPES = _Set( "object", "list", @@ -26,7 +44,8 @@ class PythonPostfixTemplatesUtils { "bool", "classmethod", "staticmethod", - "type" + "type", + "frozenset" ); /*