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" });
+ }
+
 }