Make it work with python 2 and 3. Fixes #6.

make-prefetch.py

@@ -1,12 +1,18 @@
 #!/usr/bin/env python
 
 from __future__ import print_function
+
+try:
+  from urllib.parse import urlparse
+  from urllib.request import urlopen
+except ImportError:
+  from urlparse import urlparse
+  from urllib2 import urlopen
+
 from argparse import ArgumentParser
 import hashlib
 import os
 import sys
-import urllib2
-from urlparse import urlparse
 
 usage = """make-prefetch.py [options] <file or url>
 
@@ -64,7 +70,7 @@ def hash_file(args):
   }
 
 def hash_url(args):
-  response = urllib2.urlopen(args.file)
+  response = urlopen(args.file)
 
   size = 0
   sha1 = hashlib.sha1()

test/test-cases.py

@@ -12,25 +12,25 @@ class TestHelp(unittest.TestCase):
   def test_help_short(self):
     (exitcode, stdout, stderr) = run(['-h'])
     self.assertEqual(exitcode, 0)
-    self.assertTrue('Create a prefetch statement' in stdout)
+    self.assertTrue(b'Create a prefetch statement' in stdout)
 
   def test_help_long(self):
     (exitcode, stdout, stderr) = run(['--help'])
     self.assertEqual(exitcode, 0)
-    self.assertTrue('Create a prefetch statement' in stdout)
+    self.assertTrue(b'Create a prefetch statement' in stdout)
 
 class TestURL(unittest.TestCase):
 
   def test_url(self):
     (exitcode, stdout, stderr) = run(['http://localhost:36465/hodor.txt'])
 
     expected = (
-      'prefetch '
-      'hodor.txt '
-      'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size:5 '
-      'http://localhost:36465/hodor.txt '
-      'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'prefetch '
+      b'hodor.txt '
+      b'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size:5 '
+      b'http://localhost:36465/hodor.txt '
+      b'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -41,12 +41,12 @@ def test_url_name(self):
                                       'http://localhost:36465/hodor.txt'])
 
     expected = (
-      'prefetch '
-      'bran.txt '
-      'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size:5 '
-      'http://localhost:36465/hodor.txt '
-      'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'prefetch '
+      b'bran.txt '
+      b'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size:5 '
+      b'http://localhost:36465/hodor.txt '
+      b'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -56,12 +56,12 @@ def test_url_no_name(self):
     (exitcode, stdout, stderr) = run(['http://localhost:36465/'])
 
     expected = (
-      'prefetch '
-      'REPLACEME '
-      'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size:5 '
-      'http://localhost:36465/ '
-      'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'prefetch '
+      b'REPLACEME '
+      b'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size:5 '
+      b'http://localhost:36465/ '
+      b'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -73,12 +73,12 @@ def test_file(self):
     (exitcode, stdout, stderr) = run(['hodor.txt'])
 
     expected = (
-      'prefetch '
-      'hodor.txt '
-      'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size:5 '
-      'http://REPLACEME '
-      'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'prefetch '
+      b'hodor.txt '
+      b'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size:5 '
+      b'http://REPLACEME '
+      b'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -88,11 +88,11 @@ def test_file_sha1(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-a', 'sha1'])
 
     expected = (
-      'prefetch '
-      'hodor.txt '
-      'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size:5 '
-      'http://REPLACEME'
+      b'prefetch '
+      b'hodor.txt '
+      b'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size:5 '
+      b'http://REPLACEME'
     )
 
     self.assertEqual(exitcode, 0)
@@ -102,11 +102,11 @@ def test_file_sha256(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-a', 'sha256'])
 
     expected = (
-      'prefetch '
-      'hodor.txt '
-      'size:5 '
-      'http://REPLACEME '
-      'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'prefetch '
+      b'hodor.txt '
+      b'size:5 '
+      b'http://REPLACEME '
+      b'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -116,12 +116,12 @@ def test_file_url(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-u', 'http://example.com'])
 
     expected = (
-      'prefetch '
-      'hodor.txt '
-      'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size:5 '
-      'http://example.com '
-      'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'prefetch '
+      b'hodor.txt '
+      b'sha1:c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size:5 '
+      b'http://example.com '
+      b'sha256:aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -130,7 +130,7 @@ def test_file_url(self):
   def test_file_value_sha1(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-o', 'value', '-a', 'sha1'])
 
-    expected = 'c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687'
+    expected = b'c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687'
 
     self.assertEqual(exitcode, 0)
     self.assertEqual(expected, stdout.strip())
@@ -140,7 +140,7 @@ def test_file_value_sha256(self):
                                       'sha256'])
 
     expected = (
-      'aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
+      b'aa116aa7b00b3006a1cf72219634cb1edbfad6abb1708c2e01b0409800f958b5'
     )
 
     self.assertEqual(exitcode, 0)
@@ -149,21 +149,23 @@ def test_file_value_sha256(self):
   def test_file_value_error(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-o', 'value'])
 
+    expected = b'You must specify a hash algorithm to use'
+
     self.assertEqual(exitcode, 2)
-    self.assertEqual('You must specify a hash algorithm to use', stderr.strip())
+    self.assertEqual(expected, stderr.strip())
 
   def test_file_value_davis(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-o', 'davis'])
 
     expected = (
-      'begin prefetch block\n'
-      'add prefetch item '
-      'name=hodor.txt '
-      'sha1=c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
-      'size=5 '
-      'url=http://REPLACEME\n'
-      'collect prefetch items\n'
-      'end prefetch block'
+      b'begin prefetch block\n'
+      b'add prefetch item '
+      b'name=hodor.txt '
+      b'sha1=c6447b82fbb4b8e7dbcf2d28a4d7372f5dc32687 '
+      b'size=5 '
+      b'url=http://REPLACEME\n'
+      b'collect prefetch items\n'
+      b'end prefetch block'
     )
 
     self.assertEqual(exitcode, 0)
@@ -173,7 +175,7 @@ def test_file_value_davis_error(self):
     (exitcode, stdout, stderr) = run(['hodor.txt', '-o', 'davis', '-a',
                                       'sha256'])
 
-    expected = 'Algorithm sha256 is not supported in davis downloads'
+    expected = b'Algorithm sha256 is not supported in davis downloads'
 
     self.assertEqual(exitcode, 2)
     self.assertEqual(expected, stderr.strip())