Implements xylo#172

Implemented complex types which can contain multiple types. Change how type names are retrieved, because previous approach created incorrect names for tuples, e.g.: (1,2,3) were (int, int, int) and not tuple. WIP: New type names should be tested, maybe they will provide wrong results for objects
ikopysov · Oct 30, 2020 · 87872b5 · 87872b5
1 parent 2ad42e1
commit 87872b5
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 36 deletions.
diff --git a/...e/endrullis/idea/postfixtemplates/languages/python/CustomPythonStringPostfixTemplate.java b/...e/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<String, Condition<PsiElement>> type2psiCondition =
+        new HashMap<String, Condition<PsiElement>>() {{
+            put(SpecialType.ANY.name(), e -> true);
 
-	/** Contains predefined type-to-psiCondition mappings as well as cached mappings for individual types. */
-	private static final Map<String, Condition<PsiElement>> type2psiCondition = new HashMap<String, Condition<PsiElement>>() {{
-		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<PsiElement> getCondition(final @NotNull String matchingClass, final @Nullable String conditionClass) {
-		Condition<PsiElement> psiElementCondition = type2psiCondition.get(matchingClass);
+    @NotNull
+    public static Condition<PsiElement> getCondition(final @NotNull String matchingClass,
+        final @Nullable String conditionClass) {
+        Condition<PsiElement> 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
@@ -20,6 +20,9 @@ public class PythonAnnotator implements CptLangAnnotator {
 
 	private final Map<String, Boolean> className2exists = new HashMap<String, Boolean>() {{
 		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
@@ -1,16 +1,34 @@
 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.
  *
  * @author Stefan Endrullis &lt;[email protected]&gt;
  */
 class PythonPostfixTemplatesUtils {
 
+	static final Map<String, Set<String>> COMPLEX_TYPES;
+
+	static final Set<String> ITERABLE_TYPES = _Set(
+		"list",
+		"dict",
+		"set",
+		"tuple",
+		"frozenset"
+	);
+
+	static {
+		COMPLEX_TYPES = new HashMap<>();
+		// iter
+		COMPLEX_TYPES.put("iter", ITERABLE_TYPES);
+	}
+
 	static final Set<String> PYTHON_TYPES = _Set(
 		"object",
 		"list",
@@ -26,7 +44,8 @@ class PythonPostfixTemplatesUtils {
 		"bool",
 		"classmethod",
 		"staticmethod",
-		"type"
+		"type",
+		"frozenset"
 	);
 
 	/*