[Jython-checkins] jython: Add step to itertools count function.

[frank.wierzbicki]

http://hg.python.org/jython/rev/6bc032612d2e
changeset:   6577:6bc032612d2e
user:        Frank Wierzbicki <fwierzbicki at gmail.com>
date:        Wed Apr 11 17:38:46 2012 -0700
summary:
  Add step to itertools count function.

files:
  src/org/python/modules/itertools.java |  23 +++++++++++---
  1 files changed, 18 insertions(+), 5 deletions(-)


diff --git a/src/org/python/modules/itertools.java b/src/org/python/modules/itertools.java
--- a/src/org/python/modules/itertools.java
+++ b/src/org/python/modules/itertools.java
@@ -84,22 +84,35 @@
     
 
     public static PyString __doc__count = new PyString(
-            "count([firstval]) --> count object\n\nReturn a count object whose .next() "
-                    + "method returns consecutive\nintegers starting from zero or, if specified, from firstval.");
+        "count(start=0, step=1) --> count object\n\n" +
+        "Return a count object whose .next() method returns consecutive values.\n" +
+        "  Equivalent to:\n" +
+        "\n" +
+        "      def count(firstval=0, step=1):\n" +
+        "      x = firstval\n" +
+        "      while 1:\n" +
+        "          yield x\n" +
+        "          x += step\n");
+
+    public static PyIterator count(final int init) {
+        return count(init, 1);
+    }
 
     /**
      * Creates an iterator that returns consecutive integers starting at <code>init</code>.
      */
-    public static PyIterator count(final int init) {
+    public static PyIterator count(final int init, final int step) {
         return new PyIterator() {
             int counter = init;
+            int stepper = step;
 
             public PyObject __iternext__() {
-                return new PyInteger(counter++);
+                return new PyInteger(counter+=stepper);
             }
             
             public PyString __repr__() {
-                return (PyString)(Py.newString("count(%d)").__mod__(Py.newInteger(counter)));
+                return (PyString)(Py.newString("count(%d, %d)").__mod__(new PyTuple(
+                                Py.newInteger(counter), Py.newInteger(stepper))));
             }
 
         };

-- 
Repository URL: http://hg.python.org/jython