implement /to/ as generator

Author: czheo

syntax_sugar/infix.py

@@ -1,5 +1,7 @@
 from functools import partial, reduce
+from itertools import product, islice
 from .util import flip
+from .composable import compose
 
 class infix(partial):
     def __truediv__(self, right):
@@ -12,26 +14,64 @@ def __rtruediv__(self, left):
 contains = infix(lambda lst, item: item in lst)
 pair = infix(lambda a, b: (a, b))
 join = infix(lambda lst, s: s.join(lst))
-x = infix(lambda l1, l2: list(map(lambda x, y: x * y, l1, l2)))
 hasattr = infix(hasattr)
 fmap = infix(flip(map))
 ffilter = infix(flip(filter))
 freduce = infix(flip(reduce))
+take = infix(compose(list, islice))
+
+INF = float('inf')
+
+class To:
+    def __init__(self, start, end):
+        if start /of/ int and (end /of/ int or end == INF):
+            self.type = 'number'
+            self.start = start
+            self.curr = self.start
+            self.end = end
+        elif start /of/ str and end /of/ str:
+            self.type = 'char'
+            self.start = start
+            self.curr = self.start
+            self.end = end
+        else:
+            raise TypeError('Unknown range: %s to %s' % (start, end))
+
+    def __mul__(self, rhs):
+        return product(self, rhs)
+
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        if self.type == 'number':
+            if self.curr <= self.end:
+                ret = self.curr
+                self.curr += 1
+                return ret
+            else:
+                raise StopIteration
+        elif self.type == 'char':
+            if ord(self.curr) <= ord(self.end):
+                ret = self.curr
+                self.curr = chr(ord(self.curr) + 1)
+                return ret
+            else:
+                raise StopIteration
+        else:
+            raise StopIteration
+    
+    def __str__(self):
+        if self.type == 'number':
+            return super(To, self).__str__()
+        elif self.type == 'char':
+            return ''.join(self)
+        else:
+            raise NotImplementedError
 
 @infix
 def to(start, end):
-    if start /of/ int and end /of/ int:
-        return range(start, end + 1)
-    elif start /of/ str and end /of/ str:
-        ret = ''
-        i = ord(start)
-        while i <= ord(end):
-            ret += chr(i)
-            i += 1
-        return ret
-    else:
-        raise TypeError('Unknown range: %s to %s' % (start, end))
-    
+    return To(start, end)
 
 __all__ = [
     'infix',
@@ -40,9 +80,10 @@ def to(start, end):
     'pair',
     'join',
     'to',
-    'x',
+    'INF',
     'hasattr',
     'fmap',
     'freduce',
     'ffilter',
+    'take',
 ]

tests/test_infix.py

@@ -1,7 +1,9 @@
 from syntax_sugar import *
 
 def test_int_to_int():
-    assert 1 /to/ 10 == range(1, 11)
+    assert list(1 /to/ 10) == list(range(1, 11))
 
 def test_str_to_str():
-    assert 'A' /to/ 'Z' == 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    assert str('A' /to/ 'Z') == 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+def test_take():
+    assert 1 /to/ INF /take/ 5 == [1,2,3,4,5]