toolbox/xnumerate.py at master · levis501/toolbox · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#!/usr/bin/env python3

# xnumerate.py


def isXnumberable(o):
  try:
    _ = iter(o)
    if (type(o) == type('A')) and (len(o) < 2):
      return False
    return True
  except:
    return False

class xnumerate:
  def __init__(self, subject):
    _ = iter(subject) # raises TypeError of subject isn't iterable
    self._subject = subject

  def __iter__(self):
    self._iter_stack = [iter(self._subject)]
    self._index_stack = [0]
    return self
  def __next__(self):
    while len(self._iter_stack) > 0:
      top_iter = self._iter_stack[-1]
      try:
        n = next(top_iter)
        if isXnumberable(n):
          self._iter_stack.append(iter(n))
          self._index_stack.append(0)
          continue
        i = tuple(self._index_stack)
        if len(i) == 1:
          i = i[0]
        self._index_stack[-1] += 1
        return (i, n)
      except StopIteration:
        self._iter_stack.pop()
        self._index_stack.pop()
        if (len(self._index_stack) > 0):
          self._index_stack[-1] += 1
        continue
    raise StopIteration
  def sorted(self, key = lambda x: x):
    return sorted(self, key = lambda iv: key(iv[1]))


if __name__ == '__main__':
  """Unit Testing"""
  import unittest

  class XNumerateTests(unittest.TestCase):
    def test_xnumerateExists(self):
      x = xnumerate([6])
      self.assertEqual(type(x), xnumerate)

    def test_xnumerateIsIterable(self):
      _ = iter(xnumerate([6]))

    def test_xnumerateEmptySubject(self):
      i = iter(xnumerate([]))
      self.assertRaises(StopIteration, next, i)

    def test_xnumerateArgumentNotIterable(self):
      self.assertRaises(TypeError, xnumerate, 7)

    def test_xnumerateOneItem(self):
      i = iter(xnumerate([6]))
      self.assertEqual(next(i), (0, 6))

    def test_xnumerate1x1(self):
      i = iter(xnumerate([[6]]))
      self.assertEqual(next(i), ((0, 0), 6))
      self.assertRaises(StopIteration, next, i)

    def test_xnumerate1x2(self):
      i = iter(xnumerate([[6, 7]]))
      self.assertEqual(next(i), ((0, 0), 6))
      self.assertEqual(next(i), ((0, 1), 7))
      self.assertRaises(StopIteration, next, i)

    def test_xnumerate2x2(self):
      i = iter(xnumerate([[6, 7], [9, 10]]))
      self.assertEqual(next(i), ((0, 0), 6))
      self.assertEqual(next(i), ((0, 1), 7))
      self.assertEqual(next(i), ((1, 0), 9))
      self.assertEqual(next(i), ((1, 1), 10))
      self.assertRaises(StopIteration, next, i)

    def test_xnumerateStringMixed(self):
      i = iter(xnumerate(['ab',['c', 'd']]))
      self.assertEqual(next(i), ((0, 0), 'a'))
      self.assertEqual(next(i), ((0, 1), 'b'))
      self.assertEqual(next(i), ((1, 0), 'c'))
      self.assertEqual(next(i), ((1, 1), 'd'))
      self.assertRaises(StopIteration, next, i)

    def test_xnumerateString2x2(self):
      i = iter(xnumerate([['a','b'],['y','z']]))
      self.assertEqual(next(i), ((0, 0), 'a'))
      self.assertEqual(next(i), ((0, 1), 'b'))
      self.assertEqual(next(i), ((1, 0), 'y'))
      self.assertEqual(next(i), ((1, 1), 'z'))

    def test_xnumerateSingleCharacter(self):
      i = iter(xnumerate('a'))
      self.assertEqual(next(i), (0, 'a'))

    def test_sortedOneItem(self):
      x = xnumerate([6])
      i = iter(x.sorted())
      self.assertEqual(next(i), (0, 6))

    def test_sortedTwoItems(self):
      x = xnumerate([7, 6])
      i = iter(x.sorted())
      self.assertEqual(next(i), (1, 6))

    def test_sortedWithKey(self):
      x = xnumerate((6, 7))
      i = iter(x.sorted(key = lambda x: -x))
      self.assertEqual(next(i), (1, 7))


  if __name__ == '__main__':
    unittest.main()