Implement more token features

Author: Adam Hrbac

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/contextvars/ContextBuiltins.java

@@ -124,6 +124,19 @@ Object doGetDefault(PContextVarsContext self, Object key, Object def,
 
     }
 
+    @Builtin(name = "__contains__", minNumOfPositionalArgs = 2)
+    @GenerateNodeFactory
+    public abstract static class Contains extends PythonBuiltinNode {
+        @Specialization
+        boolean doIn(PContextVarsContext self, Object key) {
+            if (key instanceof PContextVar) {
+                PContextVar var = (PContextVar) key;
+                return self.contextVarValues.lookup(var, var.getHash()) != null;
+            }
+            return false;
+        }
+    }
+
     private static Object getContextVar(PContextVarsContext self, Object key, Object def,
                     @Cached PRaiseNode raise) {
         if (key instanceof PContextVar) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/contextvars/ContextVarBuiltins.java

@@ -118,8 +118,9 @@ public abstract static class ResetNode extends PythonBinaryBuiltinNode {
         Object reset(PContextVar self, PContextVarsToken token,
                         @Cached PRaiseNode raise) {
             if (self == token.getVar()) {
+                token.use(raise);
                 PythonContext.PythonThreadState threadState = getContext().getThreadState(getLanguage());
-                if (token.getOldValue() == PContextVarsToken.MISSING) {
+                if (token.getOldValue() == null) {
                     PContextVarsContext context = threadState.getContextVarsContext();
                     context.contextVarValues = context.contextVarValues.without(self, self.getHash());
                 } else {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/contextvars/PContextVarsToken.java

@@ -40,20 +40,32 @@
  */
 package com.oracle.graal.python.builtins.objects.contextvars;
 
+import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
+import com.oracle.graal.python.nodes.ErrorMessages;
+import com.oracle.graal.python.nodes.PRaiseNode;
 import com.oracle.truffle.api.object.Shape;
 
 public class PContextVarsToken extends PythonBuiltinObject {
     public static final Object MISSING = new Object();
     private final PContextVar var;
     private final Object oldValue;
 
+    private boolean used = false;
+
     public PContextVarsToken(PContextVar var, Object oldValue, Object cls, Shape instanceShape) {
         super(cls, instanceShape);
         this.var = var;
         this.oldValue = oldValue;
     }
 
+    public void use(PRaiseNode raise) {
+        if (used) {
+            throw raise.raise(PythonBuiltinClassType.RuntimeError, ErrorMessages.TOKEN_ALREADY_USED, this);
+        }
+        used = true;
+    }
+
     public PContextVar getVar() {
         return var;
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/contextvars/TokenBuiltins.java

@@ -40,17 +40,23 @@
  */
 package com.oracle.graal.python.builtins.objects.contextvars;
 
+import static com.oracle.graal.python.util.PythonUtils.tsLiteral;
+
 import java.util.List;
 
 import com.oracle.graal.python.builtins.Builtin;
 import com.oracle.graal.python.builtins.CoreFunctions;
+import com.oracle.graal.python.builtins.Python3Core;
 import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.PythonBuiltins;
+import com.oracle.graal.python.nodes.attributes.WriteAttributeToBuiltinTypeNode;
+import com.oracle.graal.python.nodes.attributes.WriteAttributeToBuiltinTypeNodeGen;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinNode;
 import com.oracle.truffle.api.dsl.GenerateNodeFactory;
 import com.oracle.truffle.api.dsl.NodeFactory;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.strings.TruffleString;
 
 @CoreFunctions(extendClasses = PythonBuiltinClassType.ContextVarsToken)
 public class TokenBuiltins extends PythonBuiltins {
@@ -74,7 +80,17 @@ public Object doVar(PContextVarsToken self) {
     public abstract static class OldValueAttr extends PythonBuiltinNode {
         @Specialization
         public Object doOld(PContextVarsToken self) {
-            return self.getOldValue();
+            Object oldValue = self.getOldValue();
+            return oldValue == null ? PContextVarsToken.MISSING : oldValue;
         }
     }
+
+    private static final TruffleString MISSING_NAME = tsLiteral("MISSING");
+
+    @Override
+    public void postInitialize(Python3Core core) {
+        super.postInitialize(core);
+        WriteAttributeToBuiltinTypeNode addMissing = WriteAttributeToBuiltinTypeNodeGen.getUncached();
+        addMissing.execute(PythonBuiltinClassType.ContextVarsToken, MISSING_NAME, PContextVarsToken.MISSING);
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/ErrorMessages.java

@@ -1458,4 +1458,5 @@ public abstract class ErrorMessages {
     public static final TruffleString SETTRACE_NOT_IMPLEMENTED = tsLiteral("sys.settrace is only implemented for the bytecode interpreter.");
     public static final TruffleString ATTRIBUTE_VALUE_MUST_BE_BOOL = tsLiteral("attribute value type must be bool");
     public static final TruffleString HPY_UNEXPECTED_HPY_NULL = tsLiteral("unexpected HPy_NULL");
+    public static final TruffleString TOKEN_ALREADY_USED = tsLiteral("%s has already been used once");
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/object/PythonObjectFactory.java

@@ -1457,6 +1457,6 @@ public final PContextVarsContext copyContextVarsContext(PContextVarsContext orig
     }
 
     public final PContextVarsToken createContextVarsToken(PContextVar var, Object oldValue) {
-        return trace(new PContextVarsToken(var, oldValue == null ? PContextVarsToken.MISSING : oldValue, PythonBuiltinClassType.ContextVarsToken, getShape(PythonBuiltinClassType.ContextVarsToken)));
+        return trace(new PContextVarsToken(var, oldValue, PythonBuiltinClassType.ContextVarsToken, getShape(PythonBuiltinClassType.ContextVarsToken)));
     }
 }