Properly consider that attributes in enum subclasses are actually ins…

…tances of the class where they're defined.
fabioz · Oct 1, 2023 · 8d835ea · 8d835ea
1 parent 243ff88
commit 8d835ea
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 25 deletions.
diff --git a/...org.python.pydev.ast/src/org/python/pydev/ast/codecompletion/revisited/AbstractToken.java b/...org.python.pydev.ast/src/org/python/pydev/ast/codecompletion/revisited/AbstractToken.java
@@ -408,7 +408,11 @@ public static boolean isFunctionDefProperty(IToken element) {
                     for (int i = 0; i < decs.length; i++) {
                         decoratorsType dec = decs[i];
                         if (dec != null && dec.func != null) {
-                            if ("property".equals(NodeUtils.getRepresentationString(dec.func))) {
+                            final String representationString = NodeUtils.getRepresentationString(dec.func);
+                            if ("property".equals(representationString)
+                                    // The ones below are used for enums (depending on the python version).
+                                    || "DynamicClassAttribute".equals(representationString)
+                                    || "_magic_enum_attr".equals(representationString)) {
                                 return true;
                             }
                         }

diff --git a/...hon.pydev.ast/src/org/python/pydev/ast/codecompletion/revisited/modules/SourceModule.java b/...hon.pydev.ast/src/org/python/pydev/ast/codecompletion/revisited/modules/SourceModule.java
@@ -565,12 +565,12 @@ public TokensList internalCalculateGlobalTokens(ICompletionState initialState, I
                                     if (assign.targets.length == 1 && assign.targets[0] instanceof Name) {
                                         ClassDef classDef = (ClassDef) d.scope.getClassDef();
                                         if (NodeUtils.isEnum(classDef)) {
-                                            return new TokensList(new IToken[] {
-                                                    new ConcreteToken("name", "Enum name", "", "enum",
-                                                            ConcreteToken.TYPE_ATTR, manager.getNature()),
-                                                    new ConcreteToken("value", "Enum value", "", "enum",
-                                                            ConcreteToken.TYPE_ATTR, manager.getNature())
-                                            });
+                                            TokensList toks = ((SourceModule) d.module).getClassToks(
+                                                    initialState,
+                                                    manager,
+                                                    classDef);
+                                            toks.setLookingFor(LookingFor.LOOKING_FOR_INSTANCED_VARIABLE);
+                                            return toks;
                                         }
                                     }
 
@@ -604,12 +604,12 @@ public TokensList internalCalculateGlobalTokens(ICompletionState initialState, I
                                     ClassDef classDef = (ClassDef) d.scope.getClassDef();
                                     if (classDef != null) {
                                         if (NodeUtils.isEnum(classDef)) {
-                                            return new TokensList(new IToken[] {
-                                                    new ConcreteToken("name", "Enum name", "", "enum",
-                                                            ConcreteToken.TYPE_ATTR, manager.getNature()),
-                                                    new ConcreteToken("value", "Enum value", "", "enum",
-                                                            ConcreteToken.TYPE_ATTR, manager.getNature())
-                                            });
+                                            TokensList toks = ((SourceModule) d.module).getClassToks(
+                                                    initialState,
+                                                    manager,
+                                                    classDef);
+                                            toks.setLookingFor(LookingFor.LOOKING_FOR_INSTANCED_VARIABLE);
+                                            return toks;
                                         }
 
                                         FindDefinitionModelVisitor visitor = new FindDefinitionModelVisitor(

diff --git a/..._completions/org/python/pydev/ast/codecompletion/PythonCompletionWithoutBuiltinsTest.java b/..._completions/org/python/pydev/ast/codecompletion/PythonCompletionWithoutBuiltinsTest.java
@@ -3182,18 +3182,6 @@ public void testCodeCompletionFromAliasedImport() throws Exception {
         assertEquals(1, comps.length);
     }
 
-    public void testEnum() throws Exception {
-        String s;
-        s = "" +
-                "from enum import Enum\n" +
-                "\n" +
-                "class Color(Enum):\n" +
-                "    Black = '#000000'\n" +
-                "    White = '#ffffff'\n" +
-                "Color.Black.";
-        requestCompl(s, s.length(), -1, new String[] { "name", "value" });
-    }
-
     public void testCallChain() throws Exception {
         String s;
         s = "" +
@@ -3525,4 +3513,49 @@ public void testForwardDeclarationCompletion2() throws Exception {
         requestCompl(s, s.length(), -1, new String[] { "d()" });
     }
 
+    public void testEnumNameVal() throws Exception {
+        String s;
+        s = "" +
+                "from enum import Enum\n" +
+                "\n" +
+                "class Color(Enum):\n" +
+                "    Black = '#000000'\n" +
+                "    White = '#ffffff'\n" +
+                "Color.Black.";
+        requestCompl(s, s.length(), -1, new String[] { "name", "value" });
+    }
+
+    public void testEnumCustom() throws Exception {
+        String s;
+        s = "" +
+                "from enum import Enum\n"
+                + "\n"
+                + "class Signal(Enum):\n"
+                + "    INITIALISE = 1\n"
+                + "\n"
+                + "    @property\n"
+                + "    def signal_type(self) -> str:\n"
+                + "        return \"some type\"\n"
+                + "\n"
+                + "Signal.INITIALISE.";
+        requestCompl(s, s.length(), -1, new String[] { "name", "value", "signal_type" });
+    }
+
+    public void testEnumCustom2() throws Exception {
+        String s;
+        s = "" +
+                "from enum import Enum\n"
+                + "\n"
+                + "class Signal(Enum):\n"
+                + "    INITIALISE = 1\n"
+                + "\n"
+                + "    @property\n"
+                + "    def signal_type(self) -> str:\n"
+                + "        return \"some type\"\n"
+                + "\n"
+                + "X = Signal.INITIALISE\n"
+                + "X.";
+        requestCompl(s, s.length(), -1, new String[] { "name", "value", "signal_type" });
+    }
+
 }