postfix completion for 'len' function on 'Sized' objects

2026-09-27 10:03:11 +07:00 · 2019-02-28 14:14:23 +00:00
parent dfc06cb2b4
commit 847811c72e
24 changed files with 149 additions and 2 deletions
@@ -0,0 +1,2 @@
+def f(a):
+    len(a)<spot></spot>
@@ -0,0 +1,2 @@
+def f(a):
+    <spot>a</spot>.len
@@ -0,0 +1,5 @@
+<html>
+<body>
+Converts postfix call to '.len' to 'len' function invocation.
+</body>
+</html>
@@ -0,0 +1,38 @@
+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.codeInsight.postfix
+
+import com.intellij.codeInsight.template.postfix.templates.SurroundPostfixTemplateBase
+import com.intellij.lang.surroundWith.Surrounder
+import com.intellij.openapi.util.Condition
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.PyNames
+import com.jetbrains.python.psi.AccessDirection
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyExpression
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.TypeEvalContext
+import com.jetbrains.python.refactoring.surround.surrounders.expressions.PyLenExpressionStatementSurrounder
+
+class PyLenPostfixTemplate : SurroundPostfixTemplateBase("len", DESCR, PyPostfixUtils.PY_PSI_INFO,
+                                                         PyPostfixUtils.selectorAllExpressionsWithCurrentOffset(sizedFilter)) {
+
+  override fun getSurrounder(): Surrounder = PyLenExpressionStatementSurrounder()
+
+  companion object {
+    const val DESCR = "len(expr)"
+
+    val sizedFilter: Condition<PsiElement> = Condition { element ->
+      val ref = element.reference
+      if (ref?.resolve() is PyClass)
+        return@Condition false
+
+      val expression = element as PyExpression
+      val context = TypeEvalContext.codeCompletion(expression.project, expression.containingFile)
+      val type = context.getType(expression) ?: return@Condition false
+
+      val resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(context)
+      val results = type.resolveMember(PyNames.LEN, null, AccessDirection.READ, resolveContext)
+      return@Condition results?.isNotEmpty() ?: false
+    }
+  }
+}
@@ -28,7 +28,8 @@ public class PyPostfixTemplateProvider implements PostfixTemplateProvider {
                                    new PyIsNonePostfixTemplate(),
                                    new PyIsNotNonePostfixTemplate(),
                                    new PyPrintPostfixTemplate(),
-                                    new PyMainPostfixTemplate());
+                                    new PyMainPostfixTemplate(),
+                                    new PyLenPostfixTemplate());
  }

  @Override
@@ -8,11 +8,12 @@ import com.intellij.psi.PsiFile;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.refactoring.PyRefactoringUtil;
 import com.jetbrains.python.refactoring.surround.surrounders.expressions.*;
+import com.jetbrains.python.refactoring.surround.surrounders.expressions.PyLenExpressionStatementSurrounder;
 import org.jetbrains.annotations.NotNull;

 public class PyExpressionSurroundDescriptor implements SurroundDescriptor {
  private static final Surrounder[] SURROUNDERS = {new PyWithParenthesesSurrounder(), new PyIfExpressionSurrounder(),
-    new PyWhileExpressionSurrounder(), new PyIsNoneSurrounder(), new PyIsNotNoneSurrounder()};
+    new PyWhileExpressionSurrounder(), new PyIsNoneSurrounder(), new PyIsNotNoneSurrounder(), new PyLenExpressionStatementSurrounder()};

  @Override
  @NotNull
@@ -0,0 +1,31 @@
+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.refactoring.surround.surrounders.expressions
+
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.TextRange
+import com.intellij.util.IncorrectOperationException
+import com.jetbrains.python.psi.*
+
+
+class PyLenExpressionStatementSurrounder : PyExpressionSurrounder() {
+
+  @Throws(IncorrectOperationException::class)
+  override fun surroundExpression(project: Project, editor: Editor, element: PyExpression): TextRange {
+    val range = element.textRange
+    val statement = PyElementGenerator.getInstance(project).createFromText(LanguageLevel.getDefault(), PyExpressionStatement::class.java, "len(a)")
+    val callExpression = statement.expression as PyCallExpression
+    val arg = callExpression.arguments[0]
+    arg.replace(element)
+    element.replace(callExpression)
+    return TextRange.from(range.endOffset + FUNCTION_LENGTH, 0)
+  }
+
+  override fun isApplicable(expr: PyExpression) = true
+
+  override fun getTemplateDescription()= "len(expr)"
+
+  companion object {
+    const val FUNCTION_LENGTH = "len()".length
+  }
+}
@@ -0,0 +1 @@
+foo = "something".len<caret>
@@ -0,0 +1 @@
+foo = len("something")
@@ -0,0 +1 @@
+{"first": 1, "second": 2}.len<caret>
@@ -0,0 +1 @@
+len({"first": 1, "second": 2})
@@ -0,0 +1 @@
+[1, 2, 3, 4].len<caret>
@@ -0,0 +1 @@
+len([1, 2, 3, 4])
@@ -0,0 +1,5 @@
+class MyNotSizedClass:
+    pass
+
+
+MyNotSizedClass().len<caret>
@@ -0,0 +1,5 @@
+class MyNotSizedClass:
+    pass
+
+
+MyNotSizedClass().len
@@ -0,0 +1,6 @@
+class MySizedClass:
+    def __len__(self):
+        pass
+
+
+MySizedClass().len<caret>
@@ -0,0 +1,6 @@
+class MySizedClass:
+    def __len__(self):
+        pass
+
+
+MySizedClass.len<caret>
@@ -0,0 +1,6 @@
+class MySizedClass:
+    def __len__(self):
+        pass
+
+
+MySizedClass.len
@@ -0,0 +1,6 @@
+class MySizedClass:
+    def __len__(self):
+        pass
+
+
+len(MySizedClass())
@@ -0,0 +1 @@
+"something".len<caret>
@@ -0,0 +1 @@
+len("something")
@@ -0,0 +1 @@
+(1, 2, 3, 4).len<caret>
@@ -0,0 +1 @@
+len((1, 2, 3, 4))
@@ -0,0 +1,23 @@
+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.postfix
+
+class PyLenPostfixTemplateTest : PyPostfixTemplateTestCase() {
+
+  fun testString() = doTest()
+
+  fun testList() = doTest()
+
+  fun testDict() = doTest()
+
+  fun testTuple() = doTest()
+
+  fun testSized() = doTest()
+
+  fun testSizedClassObj() = doTest()
+
+  fun testAsExpr() = doTest()
+
+  fun testNotSized() = doTest()
+
+  override fun getTestDataDir() = "len/"
+}