Skip to content

Commit 9a7d4f2

Browse files
author
Adam Hrbac
committed
Use HamtIterator to make Context iterable
1 parent 39c8346 commit 9a7d4f2

File tree

7 files changed

+141
-16
lines changed

7 files changed

+141
-16
lines changed

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/Python3Core.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,6 @@
5555
import java.util.ServiceLoader;
5656
import java.util.logging.Level;
5757

58-
import com.oracle.graal.python.builtins.objects.asyncio.AsyncGenSendBuiltins;
59-
import com.oracle.graal.python.builtins.objects.asyncio.AsyncGenThrowBuiltins;
60-
import com.oracle.graal.python.builtins.objects.asyncio.AsyncGeneratorBuiltins;
6158
import org.graalvm.nativeimage.ImageInfo;
6259

6360
import com.oracle.graal.python.PythonLanguage;
@@ -202,6 +199,9 @@
202199
import com.oracle.graal.python.builtins.objects.NotImplementedBuiltins;
203200
import com.oracle.graal.python.builtins.objects.PNone;
204201
import com.oracle.graal.python.builtins.objects.array.ArrayBuiltins;
202+
import com.oracle.graal.python.builtins.objects.asyncio.AsyncGenSendBuiltins;
203+
import com.oracle.graal.python.builtins.objects.asyncio.AsyncGenThrowBuiltins;
204+
import com.oracle.graal.python.builtins.objects.asyncio.AsyncGeneratorBuiltins;
205205
import com.oracle.graal.python.builtins.objects.asyncio.CoroutineWrapperBuiltins;
206206
import com.oracle.graal.python.builtins.objects.bool.BoolBuiltins;
207207
import com.oracle.graal.python.builtins.objects.bytes.ByteArrayBuiltins;
@@ -211,6 +211,7 @@
211211
import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
212212
import com.oracle.graal.python.builtins.objects.complex.ComplexBuiltins;
213213
import com.oracle.graal.python.builtins.objects.contextvars.ContextBuiltins;
214+
import com.oracle.graal.python.builtins.objects.contextvars.ContextIteratorBuiltins;
214215
import com.oracle.graal.python.builtins.objects.contextvars.ContextVarBuiltins;
215216
import com.oracle.graal.python.builtins.objects.contextvars.TokenBuiltins;
216217
import com.oracle.graal.python.builtins.objects.deque.DequeBuiltins;
@@ -519,9 +520,12 @@ private static PythonBuiltins[] initializeBuiltins(boolean nativeAccessAllowed)
519520
new WeakRefModuleBuiltins(),
520521
new ReferenceTypeBuiltins(),
521522
new TracemallocModuleBuiltins(),
523+
// contextvars
522524
new ContextVarBuiltins(),
523525
new ContextBuiltins(),
524526
new TokenBuiltins(),
527+
new ContextIteratorBuiltins(),
528+
525529
new GenericAliasBuiltins(),
526530
new com.oracle.graal.python.builtins.objects.types.UnionTypeBuiltins(),
527531
// exceptions

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/PythonBuiltinClassType.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,8 @@ public enum PythonBuiltinClassType implements TruffleObject {
461461
ContextVarsToken("Token", J__CONTEXTVARS, Flags.PUBLIC_DERIVED_WODICT),
462462
ContextVarsContext("Context", J__CONTEXTVARS, Flags.PUBLIC_DERIVED_WODICT, CONTEXT_M_FLAGS),
463463
ContextVar("ContextVar", J__CONTEXTVARS, Flags.PUBLIC_DERIVED_WODICT),
464+
// CPython uses separate keys, values, items python types for the iterators.
465+
ContextIterator("context_iterator", J__CONTEXTVARS, Flags.PUBLIC_DERIVED_WODICT),
464466

465467
Capsule("capsule"),
466468

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242

4343
import static com.oracle.graal.python.nodes.SpecialMethodNames.J___CONTAINS__;
4444
import static com.oracle.graal.python.nodes.SpecialMethodNames.J___GETITEM__;
45+
import static com.oracle.graal.python.nodes.SpecialMethodNames.J___ITER__;
4546
import static com.oracle.graal.python.nodes.SpecialMethodNames.J___LEN__;
4647

4748
import java.util.List;
@@ -93,6 +94,42 @@ Object get(PContextVarsContext self, Object key,
9394
}
9495
}
9596

97+
@Builtin(name = J___ITER__, declaresExplicitSelf = true, minNumOfPositionalArgs = 1)
98+
@GenerateNodeFactory
99+
public abstract static class Iter extends PythonUnaryBuiltinNode {
100+
@Specialization
101+
Object iter(PContextVarsContext self) {
102+
return factory().createContextIterator(self, PContextIterator.ItemKind.KEYS);
103+
}
104+
}
105+
106+
@Builtin(name = "keys", declaresExplicitSelf = true, minNumOfPositionalArgs = 1)
107+
@GenerateNodeFactory
108+
public abstract static class Keys extends PythonUnaryBuiltinNode {
109+
@Specialization
110+
Object keys(PContextVarsContext self) {
111+
return factory().createContextIterator(self, PContextIterator.ItemKind.KEYS);
112+
}
113+
}
114+
115+
@Builtin(name = "values", declaresExplicitSelf = true, minNumOfPositionalArgs = 1)
116+
@GenerateNodeFactory
117+
public abstract static class Values extends PythonUnaryBuiltinNode {
118+
@Specialization
119+
Object values(PContextVarsContext self) {
120+
return factory().createContextIterator(self, PContextIterator.ItemKind.VALUES);
121+
}
122+
}
123+
124+
@Builtin(name = "items", declaresExplicitSelf = true, minNumOfPositionalArgs = 1)
125+
@GenerateNodeFactory
126+
public abstract static class Items extends PythonUnaryBuiltinNode {
127+
@Specialization
128+
Object items(PContextVarsContext self) {
129+
return factory().createContextIterator(self, PContextIterator.ItemKind.ITEMS);
130+
}
131+
}
132+
96133
@Builtin(name = "run", takesVarArgs = true, takesVarKeywordArgs = true, minNumOfPositionalArgs = 2, parameterNames = {"$self", "$callable"})
97134
@GenerateNodeFactory
98135
public abstract static class Run extends PythonBuiltinNode {
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package com.oracle.graal.python.builtins.objects.contextvars;
2+
3+
import static com.oracle.graal.python.nodes.SpecialMethodNames.J___ITER__;
4+
import static com.oracle.graal.python.nodes.SpecialMethodNames.J___NEXT__;
5+
6+
import java.util.List;
7+
8+
import com.oracle.graal.python.builtins.Builtin;
9+
import com.oracle.graal.python.builtins.CoreFunctions;
10+
import com.oracle.graal.python.builtins.PythonBuiltinClassType;
11+
import com.oracle.graal.python.builtins.PythonBuiltins;
12+
import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
13+
import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
14+
import com.oracle.truffle.api.dsl.GenerateNodeFactory;
15+
import com.oracle.truffle.api.dsl.NodeFactory;
16+
import com.oracle.truffle.api.dsl.Specialization;
17+
18+
@CoreFunctions(extendClasses = PythonBuiltinClassType.ContextIterator)
19+
public class ContextIteratorBuiltins extends PythonBuiltins {
20+
@Override
21+
protected List<? extends NodeFactory<? extends PythonBuiltinBaseNode>> getNodeFactories() {
22+
return ContextIteratorBuiltinsFactory.getFactories();
23+
}
24+
25+
@Builtin(name = J___ITER__, declaresExplicitSelf = true, minNumOfPositionalArgs = 1)
26+
@GenerateNodeFactory
27+
public abstract static class Iter extends PythonUnaryBuiltinNode {
28+
@Specialization
29+
static Object iter(PContextIterator self) {
30+
return self;
31+
}
32+
}
33+
34+
@Builtin(name = J___NEXT__, declaresExplicitSelf = true, minNumOfPositionalArgs = 1)
35+
@GenerateNodeFactory
36+
public abstract static class Next extends PythonUnaryBuiltinNode {
37+
@Specialization
38+
Object next(PContextIterator self) {
39+
Object next = self.next(factory());
40+
if (next == null) {
41+
throw raiseStopIteration();
42+
} else {
43+
return next;
44+
}
45+
}
46+
}
47+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package com.oracle.graal.python.builtins.objects.contextvars;
2+
3+
import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
4+
import com.oracle.graal.python.runtime.object.PythonObjectFactory;
5+
import com.oracle.truffle.api.CompilerDirectives;
6+
import com.oracle.truffle.api.object.Shape;
7+
8+
public final class PContextIterator extends PythonBuiltinObject {
9+
public enum ItemKind {
10+
KEYS,
11+
VALUES,
12+
ITEMS;
13+
14+
public Object apply(Hamt.Entry item, PythonObjectFactory factory) {
15+
switch (this) {
16+
case KEYS:
17+
return item.key;
18+
case VALUES:
19+
return item.value;
20+
case ITEMS:
21+
return factory.createTuple(new Object[]{item.key, item.value});
22+
default:
23+
throw CompilerDirectives.shouldNotReachHere("null ItemKind in PHamtIterator");
24+
}
25+
}
26+
}
27+
28+
private final ItemKind kind;
29+
private final HamtIterator it;
30+
31+
public PContextIterator(Object cls, Shape instanceShape, PContextVarsContext ctx, ItemKind kind) {
32+
super(cls, instanceShape);
33+
this.it = new HamtIterator(ctx.contextVarValues);
34+
this.kind = kind;
35+
36+
}
37+
38+
// can return null
39+
public Object next(PythonObjectFactory factory) {
40+
Hamt.Entry item = it.next();
41+
return item == null ? null : kind.apply(item, factory);
42+
}
43+
}

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

Lines changed: 0 additions & 13 deletions
This file was deleted.

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@
9191
import com.oracle.graal.python.builtins.objects.common.HashingStorageNodes;
9292
import com.oracle.graal.python.builtins.objects.common.PHashingCollection;
9393
import com.oracle.graal.python.builtins.objects.complex.PComplex;
94+
import com.oracle.graal.python.builtins.objects.contextvars.PContextIterator;
9495
import com.oracle.graal.python.builtins.objects.contextvars.PContextVar;
9596
import com.oracle.graal.python.builtins.objects.contextvars.PContextVarsContext;
9697
import com.oracle.graal.python.builtins.objects.contextvars.PContextVarsToken;
@@ -1463,6 +1464,10 @@ public final PContextVarsContext createContextVarsContext() {
14631464
return trace(new PContextVarsContext(PythonBuiltinClassType.ContextVarsContext, getShape(PythonBuiltinClassType.ContextVarsContext)));
14641465
}
14651466

1467+
public final PContextIterator createContextIterator(PContextVarsContext ctx, PContextIterator.ItemKind kind) {
1468+
return trace(new PContextIterator(PythonBuiltinClassType.ContextIterator, getShape(PythonBuiltinClassType.ContextIterator), ctx, kind));
1469+
}
1470+
14661471
public final PContextVarsContext copyContextVarsContext(PContextVarsContext original) {
14671472
return trace(new PContextVarsContext(original, PythonBuiltinClassType.ContextVarsContext, getShape(PythonBuiltinClassType.ContextVarsContext)));
14681473
}

0 commit comments

Comments
 (0)