[GR-25388] Support start and end arguments in the bytes.index method and implement itertools.tee.

PullRequest: graalpython/1176

Author: steve-s

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017, 2018, Oracle and/or its affiliates.
+ * Copyright (c) 2017, 2020, Oracle and/or its affiliates.
  * Copyright (c) 2013, Regents of the University of California
  *
  * All rights reserved.
@@ -25,37 +25,19 @@
  */
 package com.oracle.graal.python.builtins.modules;
 
-import static com.oracle.graal.python.runtime.exception.PythonErrorType.NotImplementedError;
-
+import java.util.Collections;
 import java.util.List;
 
-import com.oracle.graal.python.builtins.Builtin;
 import com.oracle.graal.python.builtins.CoreFunctions;
 import com.oracle.graal.python.builtins.PythonBuiltins;
-import com.oracle.graal.python.builtins.objects.PNone;
 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;
 
 @CoreFunctions(defineModule = "itertools")
 public final class ItertoolsModuleBuiltins extends PythonBuiltins {
 
     @Override
     protected List<? extends NodeFactory<? extends PythonBuiltinBaseNode>> getNodeFactories() {
-        return ItertoolsModuleBuiltinsFactory.getFactories();
-    }
-
-    // tee(iterable, n=2)
-    @Builtin(name = "tee", minNumOfPositionalArgs = 1, maxNumOfPositionalArgs = 2)
-    @GenerateNodeFactory
-    public abstract static class TeeNode extends PythonBuiltinNode {
-        @SuppressWarnings("unused")
-        @Specialization(guards = "isNoValue(n)")
-        Object tee(Object iterable, PNone n) {
-            throw raise(NotImplementedError, "tee");
-        }
+        return Collections.emptyList();
     }
-
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/bytes/BytesBuiltins.java

@@ -62,6 +62,7 @@
 import com.oracle.graal.python.builtins.objects.PNotImplemented;
 import com.oracle.graal.python.builtins.objects.bytes.BytesBuiltinsFactory.BytesLikeNoGeneralizationNodeGen;
 import com.oracle.graal.python.builtins.objects.common.IndexNodes.NormalizeIndexNode;
+import com.oracle.graal.python.builtins.objects.common.SequenceNodes;
 import com.oracle.graal.python.builtins.objects.common.SequenceNodes.GetObjectArrayNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceNodes.GetSequenceStorageNode;
 import com.oracle.graal.python.builtins.objects.common.SequenceNodesFactory.GetObjectArrayNodeGen;
@@ -853,23 +854,42 @@ protected boolean doIt(byte[] bytes, byte[] suffix, int start, int end) {
 
     // bytes.index(x)
     // bytearray.index(x)
-    @Builtin(name = "index", minNumOfPositionalArgs = 2)
+    @Builtin(name = "index", minNumOfPositionalArgs = 2, maxNumOfPositionalArgs = 4)
     @GenerateNodeFactory
-    public abstract static class ByteArrayIndexNode extends PythonBuiltinNode {
-        @Child private SequenceStorageNodes.LenNode lenNode;
-
+    public abstract static class ByteArrayIndexNode extends PythonQuaternaryBuiltinNode {
         @Specialization
-        public int index(VirtualFrame frame, PIBytesLike byteArray, Object arg,
-                        @Cached("create()") BytesNodes.FindNode findNode) {
-            return findNode.execute(frame, byteArray, arg, 0, getLength(byteArray.getSequenceStorage()));
+        int index(VirtualFrame frame, PIBytesLike byteArray, Object arg, @SuppressWarnings("unused") PNone start, @SuppressWarnings("unused") PNone end,
+                        @Shared("len") @Cached SequenceStorageNodes.LenNode lenNode,
+                        @Shared("storage") @Cached SequenceNodes.GetSequenceStorageNode getStorageNode,
+                        @Shared("findNode") @Cached("create()") BytesNodes.FindNode findNode) {
+            return checkResult(findNode.execute(frame, byteArray, arg, 0, lenNode.execute(getStorageNode.execute(byteArray))));
         }
 
-        private int getLength(SequenceStorage s) {
-            if (lenNode == null) {
-                CompilerDirectives.transferToInterpreterAndInvalidate();
-                lenNode = insert(SequenceStorageNodes.LenNode.create());
+        @Specialization(guards = {"!isPNone(start)"})
+        int indexWithStart(VirtualFrame frame, PIBytesLike byteArray, Object arg, Object start, @SuppressWarnings("unused") PNone end,
+                        @Shared("storage") @Cached SequenceNodes.GetSequenceStorageNode getStorageNode,
+                        @Shared("len") @Cached SequenceStorageNodes.LenNode lenNode,
+                        @Shared("findNode") @Cached("create()") BytesNodes.FindNode findNode) {
+            return checkResult(findNode.execute(frame, byteArray, arg, start, lenNode.execute(getStorageNode.execute(byteArray))));
+        }
+
+        @Specialization(guards = {"!isPNone(end)"})
+        int indexWithEnd(VirtualFrame frame, PIBytesLike byteArray, Object arg, @SuppressWarnings("unused") PNone start, Object end,
+                        @Shared("findNode") @Cached("create()") BytesNodes.FindNode findNode) {
+            return checkResult(checkResult(findNode.execute(frame, byteArray, arg, 0, end)));
+        }
+
+        @Specialization(guards = {"!isPNone(start)", "!isPNone(end)"})
+        int indexWithStartEnd(VirtualFrame frame, PIBytesLike byteArray, Object arg, Object start, Object end,
+                        @Shared("findNode") @Cached("create()") BytesNodes.FindNode findNode) {
+            return checkResult(findNode.execute(frame, byteArray, arg, start, end));
+        }
+
+        private int checkResult(int result) {
+            if (result == -1) {
+                throw raise(PythonBuiltinClassType.ValueError, ErrorMessages.SUBSECTION_NOT_FOUND);
             }
-            return lenNode.execute(s);
+            return result;
         }
     }
 

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

@@ -473,6 +473,7 @@ public abstract class ErrorMessages {
     public static final String STRING_ARG_WO_ENCODING = "string argument without an encoding";
     public static final String STRING_INDEX_OUT_OF_RANGE = "IndexError: string index out of range";
     public static final String SUBSTRING_NOT_FOUND = "substring not found";
+    public static final String SUBSECTION_NOT_FOUND = "subsection not found";
     public static final String SUPER_OBJ_MUST_BE_INST_SUB_OR_TYPE = "super(type, obj): obj must be an instance or subtype of type";
     public static final String TAKES_D_OR_D_ARGS = "%s takes %d or %d arguments";
     public static final String TAKES_D_POS_ARG_S_BUT_D_POS_ARG_S = "%s() takes %d positional argument%s but %d positional argument%s (and %d keyword-only argument%s) were given%s";

graalpython/lib-graalpython/itertools.py

@@ -858,3 +858,34 @@ def __next__(self):
     @__graalpython__.builtin_method
     def __reduce__(self):
         return (type(self), (self.data, self.selectors))
+
+
+def tee(iterable, n=2):
+    import collections
+    class _tee:
+        @__graalpython__.builtin_method
+        def __init__(self, it, deque, deques):
+            self.it = it
+            self.deque = deque
+            self.deques = deques
+
+        @__graalpython__.builtin_method
+        def __iter__(self):
+            return self
+
+        @__graalpython__.builtin_method
+        def __next__(self):
+            if not self.deque:
+                newval = next(self.it)
+                for d in self.deques:
+                    d.append(newval)
+            return self.deque.popleft()
+
+    if not isinstance(n, int):
+        raise TypeError()
+    if n < 0:
+        raise ValueError("n must be >=0")
+
+    deques = [collections.deque() for i in range(n)]
+    it = iter(iterable)
+    return tuple(_tee(it, d, deques) for d in deques)