Skip to content

Commit a8c0d42

Browse files
tomasstupkatomasstupka
committed
[GR-33928] Intrinsify itertools.
PullRequest: graalpython/1981
2 parents 6451a25 + 3740a11 commit a8c0d42

24 files changed

+1403
-244
lines changed

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,10 @@
195195
import com.oracle.graal.python.builtins.objects.iterator.IteratorBuiltins;
196196
import com.oracle.graal.python.builtins.objects.iterator.PZipBuiltins;
197197
import com.oracle.graal.python.builtins.objects.iterator.SentinelIteratorBuiltins;
198+
import com.oracle.graal.python.builtins.objects.itertools.ChainBuiltins;
199+
import com.oracle.graal.python.builtins.objects.itertools.RepeatBuiltins;
200+
import com.oracle.graal.python.builtins.objects.itertools.TeeBuiltins;
201+
import com.oracle.graal.python.builtins.objects.itertools.TeeDataObjectBuiltins;
198202
import com.oracle.graal.python.builtins.objects.list.ListBuiltins;
199203
import com.oracle.graal.python.builtins.objects.map.MapBuiltins;
200204
import com.oracle.graal.python.builtins.objects.mappingproxy.MappingproxyBuiltins;
@@ -497,6 +501,12 @@ private static PythonBuiltins[] initializeBuiltins(boolean nativeAccessAllowed)
497501
new ZipImporterBuiltins(),
498502
new ZipImportModuleBuiltins(),
499503

504+
// itertools
505+
new ChainBuiltins(),
506+
new RepeatBuiltins(),
507+
new TeeBuiltins(),
508+
new TeeDataObjectBuiltins(),
509+
500510
// zlib
501511
new ZLibModuleBuiltins(),
502512
new ZlibCompressBuiltins(),

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,12 @@ public enum PythonBuiltinClassType implements TruffleObject {
205205
PSSLSocket("_SSLSocket", "_ssl"),
206206
PMemoryBIO("MemoryBIO", "_ssl"),
207207

208+
// itertools
209+
PTee("_tee", "itertools", Flags.PUBLIC_DERIVED_WODICT),
210+
PTeeDataObject("_tee_dataobject", "itertools", Flags.PUBLIC_DERIVED_WODICT),
211+
PRepeat("repeat", "itertools"),
212+
PChain("chain", "itertools"),
213+
208214
// json
209215
JSONScanner("Scanner", "_json", Flags.PUBLIC_BASE_WODICT),
210216
JSONEncoder("Encoder", "_json", Flags.PUBLIC_BASE_WODICT),

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/ItertoolsModuleBuiltins.java

Lines changed: 135 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,19 +25,152 @@
2525
*/
2626
package com.oracle.graal.python.builtins.modules;
2727

28-
import java.util.Collections;
28+
import static com.oracle.graal.python.nodes.ErrorMessages.S_MUST_BE_S;
29+
import static com.oracle.graal.python.builtins.PythonBuiltinClassType.ValueError;
30+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__COPY__;
31+
32+
import com.oracle.graal.python.annotations.ArgumentClinic;
33+
import com.oracle.graal.python.builtins.Builtin;
2934
import java.util.List;
3035

3136
import com.oracle.graal.python.builtins.CoreFunctions;
37+
import com.oracle.graal.python.builtins.PythonBuiltinClassType;
3238
import com.oracle.graal.python.builtins.PythonBuiltins;
39+
import com.oracle.graal.python.builtins.modules.BuiltinFunctions.IterNode;
40+
import com.oracle.graal.python.builtins.objects.PNone;
41+
import com.oracle.graal.python.builtins.objects.function.PKeyword;
42+
import com.oracle.graal.python.builtins.objects.itertools.PChain;
43+
import com.oracle.graal.python.builtins.objects.itertools.PRepeat;
44+
import com.oracle.graal.python.builtins.objects.itertools.PTeeDataObject;
45+
import com.oracle.graal.python.builtins.objects.type.TypeNodes;
46+
import com.oracle.graal.python.lib.PyCallableCheckNode;
47+
import com.oracle.graal.python.lib.PyObjectLookupAttr;
48+
import com.oracle.graal.python.nodes.ErrorMessages;
49+
import com.oracle.graal.python.nodes.call.special.CallVarargsMethodNode;
3350
import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
51+
import com.oracle.graal.python.nodes.function.builtins.PythonBinaryClinicBuiltinNode;
52+
import com.oracle.graal.python.nodes.function.builtins.PythonVarargsBuiltinNode;
53+
import com.oracle.graal.python.nodes.function.builtins.clinic.ArgumentClinicProvider;
54+
import static com.oracle.graal.python.runtime.exception.PythonErrorType.TypeError;
55+
import com.oracle.graal.python.util.PythonUtils;
56+
import com.oracle.truffle.api.dsl.Cached;
57+
import com.oracle.truffle.api.dsl.Fallback;
58+
import com.oracle.truffle.api.dsl.GenerateNodeFactory;
3459
import com.oracle.truffle.api.dsl.NodeFactory;
60+
import com.oracle.truffle.api.dsl.Specialization;
61+
import com.oracle.truffle.api.frame.VirtualFrame;
62+
import com.oracle.truffle.api.profiles.BranchProfile;
3563

3664
@CoreFunctions(defineModule = "itertools")
3765
public final class ItertoolsModuleBuiltins extends PythonBuiltins {
3866

3967
@Override
4068
protected List<? extends NodeFactory<? extends PythonBuiltinBaseNode>> getNodeFactories() {
41-
return Collections.emptyList();
69+
return ItertoolsModuleBuiltinsFactory.getFactories();
70+
}
71+
72+
@Builtin(name = "tee", minNumOfPositionalArgs = 1, parameterNames = {"iterable", "n"})
73+
@ArgumentClinic(name = "n", conversion = ArgumentClinic.ClinicConversion.Int, defaultValue = "2")
74+
@GenerateNodeFactory
75+
public abstract static class TeeNode extends PythonBinaryClinicBuiltinNode {
76+
@Override
77+
protected ArgumentClinicProvider getArgumentClinic() {
78+
return ItertoolsModuleBuiltinsClinicProviders.TeeNodeClinicProviderGen.INSTANCE;
79+
}
80+
81+
@SuppressWarnings("unused")
82+
@Specialization(guards = "n < 0")
83+
protected Object negativeN(Object iterable, int n) {
84+
throw raise(ValueError, S_MUST_BE_S, "n", ">=0");
85+
}
86+
87+
@SuppressWarnings("unused")
88+
@Specialization(guards = "n == 0")
89+
protected Object zeroN(Object iterable, int n) {
90+
return factory().createTuple(PythonUtils.EMPTY_OBJECT_ARRAY);
91+
}
92+
93+
@Specialization(guards = "n > 0")
94+
protected Object tee(VirtualFrame frame, Object iterable, int n,
95+
@Cached IterNode iterNode,
96+
@Cached PyObjectLookupAttr getAttrNode,
97+
@Cached PyCallableCheckNode callableCheckNode,
98+
@Cached CallVarargsMethodNode callNode,
99+
@Cached BranchProfile notCallableProfile) {
100+
Object it = iterNode.execute(frame, iterable, PNone.NO_VALUE);
101+
Object copyCallable = getAttrNode.execute(frame, it, __COPY__);
102+
if (!callableCheckNode.execute(copyCallable)) {
103+
notCallableProfile.enter();
104+
// as in Tee.__NEW__()
105+
PTeeDataObject dataObj = factory().createTeeDataObject(it);
106+
it = factory().createTee(dataObj, 0);
107+
}
108+
109+
// return tuple([it] + [it.__copy__() for i in range(1, n)])
110+
Object[] tupleObjs = new Object[n];
111+
tupleObjs[0] = it;
112+
113+
copyCallable = getAttrNode.execute(frame, it, __COPY__);
114+
for (int i = 1; i < n; i++) {
115+
tupleObjs[i] = callNode.execute(frame, copyCallable, PythonUtils.EMPTY_OBJECT_ARRAY, PKeyword.EMPTY_KEYWORDS);
116+
}
117+
return factory().createTuple(tupleObjs);
118+
}
119+
42120
}
121+
122+
@Builtin(name = "_tee_dataobject", minNumOfPositionalArgs = 1, takesVarArgs = true, takesVarKeywordArgs = true, constructsClass = PythonBuiltinClassType.PTeeDataObject)
123+
@GenerateNodeFactory
124+
public abstract static class TeeDataObjectNode extends PythonVarargsBuiltinNode {
125+
@SuppressWarnings("unused")
126+
@Specialization(guards = "isTypeNode.execute(cls)", limit = "1")
127+
protected PTeeDataObject construct(Object cls, Object[] arguments, PKeyword[] keywords,
128+
@SuppressWarnings("unused") @Cached TypeNodes.IsTypeNode isTypeNode) {
129+
return factory().createTeeDataObject();
130+
}
131+
132+
@Fallback
133+
@SuppressWarnings("unused")
134+
protected Object notype(Object cls, Object[] arguments, PKeyword[] keywords,
135+
@SuppressWarnings("unused") @Cached TypeNodes.IsTypeNode isTypeNode) {
136+
throw raise(TypeError, ErrorMessages.IS_NOT_TYPE_OBJ, "'cls'", cls);
137+
}
138+
}
139+
140+
@Builtin(name = "repeat", minNumOfPositionalArgs = 1, takesVarArgs = true, takesVarKeywordArgs = true, constructsClass = PythonBuiltinClassType.PRepeat, doc = "repeat(object [,times]) -> create an iterator which returns the object\n" +
141+
"for the specified number of times. If not specified, returns the object\nendlessly.")
142+
@GenerateNodeFactory
143+
public abstract static class RepeatNode extends PythonVarargsBuiltinNode {
144+
@SuppressWarnings("unused")
145+
@Specialization(guards = "isTypeNode.execute(cls)", limit = "1")
146+
protected PRepeat construct(Object cls, Object[] arguments, PKeyword[] keywords,
147+
@SuppressWarnings("unused") @Cached TypeNodes.IsTypeNode isTypeNode) {
148+
return factory().createRepeat();
149+
}
150+
151+
@Fallback
152+
@SuppressWarnings("unused")
153+
protected Object notype(Object cls, Object[] arguments, PKeyword[] keywords) {
154+
throw raise(TypeError, ErrorMessages.IS_NOT_TYPE_OBJ, "'cls'", cls);
155+
}
156+
}
157+
158+
@Builtin(name = "chain", minNumOfPositionalArgs = 1, takesVarArgs = true, takesVarKeywordArgs = true, constructsClass = PythonBuiltinClassType.PChain, doc = "Return a chain object whose .__next__() method returns elements from the " +
159+
"first iterable until it is exhausted, then elements from the next iterable, until all of the iterables are exhausted.")
160+
@GenerateNodeFactory
161+
public abstract static class ChainNode extends PythonVarargsBuiltinNode {
162+
@SuppressWarnings("unused")
163+
@Specialization(guards = "isTypeNode.execute(cls)", limit = "1")
164+
protected PChain construct(Object cls, Object[] arguments, PKeyword[] keywords,
165+
@SuppressWarnings("unused") @Cached TypeNodes.IsTypeNode isTypeNode) {
166+
return factory().createChain();
167+
}
168+
169+
@Fallback
170+
@SuppressWarnings("unused")
171+
protected Object notype(Object cls, Object[] arguments, PKeyword[] keywords) {
172+
throw raise(TypeError, ErrorMessages.IS_NOT_TYPE_OBJ, "'cls'", cls);
173+
}
174+
}
175+
43176
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/WarningsModuleBuiltins.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -988,7 +988,7 @@ Object doWarn(VirtualFrame frame, PythonModule mod, Object message, Object categ
988988
try {
989989
lineno = castLong.execute(ln);
990990
} catch (CannotCastException e) {
991-
throw raise(PythonBuiltinClassType.TypeError, ErrorMessages.INTEGER_EXPECTED_GOT_FLOAT);
991+
throw raise(PythonBuiltinClassType.TypeError, ErrorMessages.S_EXPECTED_GOT_P, "integer", "float");
992992
}
993993
PDict globalsDict;
994994
if (globals instanceof PNone) {

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/zlib/ZlibCompressBuiltins.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import static com.oracle.graal.python.builtins.modules.zlib.ZLibModuleBuiltins.Z_NO_FLUSH;
4747
import static com.oracle.graal.python.builtins.modules.zlib.ZlibNodes.Z_OK;
4848
import static com.oracle.graal.python.builtins.modules.zlib.ZlibNodes.Z_STREAM_ERROR;
49+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__COPY__;
4950
import static com.oracle.graal.python.runtime.exception.PythonErrorType.ValueError;
5051
import static com.oracle.graal.python.runtime.exception.PythonErrorType.ZLibError;
5152

@@ -184,7 +185,7 @@ Object doit(ZLibCompObject self,
184185
}
185186
}
186187

187-
@Builtin(name = "__copy__", minNumOfPositionalArgs = 1)
188+
@Builtin(name = __COPY__, minNumOfPositionalArgs = 1)
188189
@GenerateNodeFactory
189190
abstract static class UndescoreCopyNode extends CopyNode {
190191
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/zlib/ZlibDecompressBuiltins.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import static com.oracle.graal.python.builtins.modules.zlib.ZlibNodes.Z_STREAM_ERROR;
4848
import static com.oracle.graal.python.nodes.ErrorMessages.INCONSISTENT_STREAM_STATE;
4949
import static com.oracle.graal.python.nodes.ErrorMessages.S_MUST_BE_GREATER_THAN_ZERO;
50+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__COPY__;
5051
import static com.oracle.graal.python.runtime.exception.PythonErrorType.ValueError;
5152
import static com.oracle.graal.python.runtime.exception.PythonErrorType.ZLibError;
5253

@@ -199,7 +200,7 @@ Object doit(ZLibCompObject self,
199200
}
200201
}
201202

202-
@Builtin(name = "__copy__", minNumOfPositionalArgs = 1)
203+
@Builtin(name = __COPY__, minNumOfPositionalArgs = 1)
203204
@GenerateNodeFactory
204205
abstract static class UndescoreCopyNode extends CopyNode {
205206
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/common/SequenceStorageNodes.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -732,7 +732,7 @@ private static Object verifyResult(VerifyNativeItemNode verifyNativeItemNode, PR
732732

733733
@GenerateUncached
734734
@ImportStatic({ListStorageType.class, SequenceStorageBaseNode.class})
735-
abstract static class GetItemSliceNode extends Node {
735+
public abstract static class GetItemSliceNode extends Node {
736736

737737
public abstract SequenceStorage execute(SequenceStorage s, int start, int stop, int step, int length);
738738

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/deque/DequeBuiltins.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
import static com.oracle.graal.python.nodes.SpecialAttributeNames.__DICT__;
5151
import static com.oracle.graal.python.nodes.SpecialMethodNames.__ADD__;
5252
import static com.oracle.graal.python.nodes.SpecialMethodNames.__CONTAINS__;
53+
import static com.oracle.graal.python.nodes.SpecialMethodNames.__COPY__;
5354
import static com.oracle.graal.python.nodes.SpecialMethodNames.__DELITEM__;
5455
import static com.oracle.graal.python.nodes.SpecialMethodNames.__EQ__;
5556
import static com.oracle.graal.python.nodes.SpecialMethodNames.__GETITEM__;
@@ -267,7 +268,7 @@ static PNone doGeneric(PDeque self) {
267268
}
268269

269270
// deque.copy()
270-
@Builtin(name = "__copy__", minNumOfPositionalArgs = 1)
271+
@Builtin(name = __COPY__, minNumOfPositionalArgs = 1)
271272
@Builtin(name = "copy", minNumOfPositionalArgs = 1)
272273
@GenerateNodeFactory
273274
public abstract static class DequeCopyNode extends PythonUnaryBuiltinNode {

0 commit comments

Comments
 (0)