Merge remote-tracking branch 'duffer/pathod-lang-generators'

3 files changed, 41 insertions, 31 deletions

diff --git a/.travis.yml b/.travis.yml
index 9d4dea88..5ec8b3bd 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -22,9 +22,9 @@ matrix:
       git:
         depth: 9999999
     - python: 3.5
-      env: SCOPE="netlib ./test/mitmproxy/script ./test/pathod/test_utils.py ./test/pathod/test_log.py"
+      env: SCOPE="netlib ./test/mitmproxy/script ./test/pathod/test_utils.py ./test/pathod/test_log.py ./test/pathod/test_language_generators.py"
     - python: 3.5
-      env: SCOPE="netlib ./test/mitmproxy/script ./test/pathod/test_utils.py ./test/pathod/test_log.py" NO_ALPN=1
+      env: SCOPE="netlib ./test/mitmproxy/script ./test/pathod/test_utils.py ./test/pathod/test_log.py ./test/pathod/test_language_generators.py" NO_ALPN=1
     - python: 2.7
       env: DOCS=1
       script: 'cd docs && make html'
diff --git a/pathod/language/generators.py b/pathod/language/generators.py
index a17e7052..20bb7ae1 100644
--- a/pathod/language/generators.py
+++ b/pathod/language/generators.py
@@ -2,17 +2,19 @@ import string
 import random
 import mmap
 
+import six
+
 DATATYPES = dict(
-    ascii_letters=string.ascii_letters,
-    ascii_lowercase=string.ascii_lowercase,
-    ascii_uppercase=string.ascii_uppercase,
-    digits=string.digits,
-    hexdigits=string.hexdigits,
-    octdigits=string.octdigits,
-    punctuation=string.punctuation,
-    whitespace=string.whitespace,
-    ascii=string.printable,
-    bytes="".join(chr(i) for i in range(256))
+    ascii_letters=string.ascii_letters.encode(),
+    ascii_lowercase=string.ascii_lowercase.encode(),
+    ascii_uppercase=string.ascii_uppercase.encode(),
+    digits=string.digits.encode(),
+    hexdigits=string.hexdigits.encode(),
+    octdigits=string.octdigits.encode(),
+    punctuation=string.punctuation.encode(),
+    whitespace=string.whitespace.encode(),
+    ascii=string.printable.encode(),
+    bytes=bytes(bytearray(range(256)))
 )
 
 
@@ -37,14 +39,21 @@ class TransformGenerator(object):
         d = self.gen.__getitem__(x)
         return self.transform(x, d)
 
-    def __getslice__(self, a, b):
-        d = self.gen.__getslice__(a, b)
-        return self.transform(a, d)
-
     def __repr__(self):
         return "'transform(%s)'" % self.gen
 
 
+def rand_byte(chars):
+    """
+        Return a random character as byte from a charset.
+    """
+    # bytearray has consistent behaviour on both Python 2 and 3
+    # while bytes does not
+    if six.PY2:
+        return random.choice(chars)
+    return bytes([random.choice(chars)])
+
+
 class RandomGenerator(object):
 
     def __init__(self, dtype, length):
@@ -55,12 +64,10 @@ class RandomGenerator(object):
         return self.length
 
     def __getitem__(self, x):
-        return random.choice(DATATYPES[self.dtype])
-
-    def __getslice__(self, a, b):
-        b = min(b, self.length)
         chars = DATATYPES[self.dtype]
-        return "".join(random.choice(chars) for x in range(a, b))
+        if isinstance(x, slice):
+            return b"".join(rand_byte(chars) for _ in range(*x.indices(self.length)))
+        return rand_byte(chars)
 
     def __repr__(self):
         return "%s random from %s" % (self.length, self.dtype)
@@ -70,17 +77,17 @@ class FileGenerator(object):
 
     def __init__(self, path):
         self.path = path
-        self.fp = file(path, "rb")
+        self.fp = open(path, "rb")
         self.map = mmap.mmap(self.fp.fileno(), 0, access=mmap.ACCESS_READ)
 
     def __len__(self):
         return len(self.map)
 
     def __getitem__(self, x):
-        return self.map.__getitem__(x)
-
-    def __getslice__(self, a, b):
-        return self.map.__getslice__(a, b)
+        if isinstance(x, slice):
+            return self.map.__getitem__(x)
+        # A slice of length 1 returns a byte object (not an integer)
+        return self.map.__getitem__(slice(x, x+1 or self.map.size()))
 
     def __repr__(self):
         return "<%s" % self.path
diff --git a/test/pathod/test_language_generators.py b/test/pathod/test_language_generators.py
index 0fceae85..51f55991 100644
--- a/test/pathod/test_language_generators.py
+++ b/test/pathod/test_language_generators.py
@@ -7,24 +7,27 @@ import tutils
 def test_randomgenerator():
     g = generators.RandomGenerator("bytes", 100)
     assert repr(g)
+    assert g[0]
+    assert len(g[0]) == 1
     assert len(g[:10]) == 10
     assert len(g[1:10]) == 9
     assert len(g[:1000]) == 100
     assert len(g[1000:1001]) == 0
-    assert g[0]
 
 
 def test_filegenerator():
     with tutils.tmpdir() as t:
         path = os.path.join(t, "foo")
         f = open(path, "wb")
-        f.write("x" * 10000)
+        f.write(b"x" * 10000)
         f.close()
         g = generators.FileGenerator(path)
         assert len(g) == 10000
-        assert g[0] == "x"
-        assert g[-1] == "x"
-        assert g[0:5] == "xxxxx"
+        assert g[0] == b"x"
+        assert g[-1] == b"x"
+        assert g[0:5] == b"xxxxx"
+        assert len(g[1:10]) == 9
+        assert len(g[10000:10001]) == 0
         assert repr(g)
         # remove all references to FileGenerator instance to close the file
         # handle.