Basic implementation of itertools.tee

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/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)